gmakstutis/graphiti
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add text episode type (#46)

Browse source 
Add a new `text` episode type and update the `extract_nodes` function to handle it.

* **EpisodeType Enum:**
  - Add `text` to the `EpisodeType` enum in `graphiti_core/nodes.py`.
  - Update the `from_str` method to handle the `text` episode type.

* **extract_nodes Function:**
  - Update the `extract_nodes` function in `graphiti_core/utils/maintenance/node_operations.py` to handle the `text` episode type.
  - Use the `message` type prompt for both `message` and `text` episodes.

---

For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/getzep/graphiti?shareId=XXXX-XXXX-XXXX-XXXX).
This commit is contained in:
Daniel Chalef 2024-08-26 15:51:13 -07:00 • committed by GitHub
parent 2d01e5d7b7
commit a6d63f0c0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
graphiti_core/nodes.py
View file

@ -46,10 +46,13 @@ class EpisodeType(Enum):
or "assistant: I'm doing well, thank you for asking." or "assistant: I'm doing well, thank you for asking."
json : str json : str
Represents an episode containing a JSON string object with structured data. Represents an episode containing a JSON string object with structured data.
text : str
Represents a plain text episode.
""" """
message = 'message' message = 'message'
json = 'json' json = 'json'
text = 'text'
@staticmethod @staticmethod
def from_str(episode_type: str): def from_str(episode_type: str):
@ -57,6 +60,8 @@ class EpisodeType(Enum):
return EpisodeType.message return EpisodeType.message
if episode_type == 'json': if episode_type == 'json':
return EpisodeType.json return EpisodeType.json
if episode_type == 'text':
return EpisodeType.text
logger.error(f'Episode type: {episode_type} not implemented') logger.error(f'Episode type: {episode_type} not implemented')
raise NotImplementedError raise NotImplementedError

2
graphiti_core/utils/maintenance/node_operations.py
View file

@ -72,7 +72,7 @@ async def extract_nodes(
) -> list[EntityNode]: ) -> list[EntityNode]:
start = time() start = time()
extracted_node_data: list[dict[str, Any]] = [] extracted_node_data: list[dict[str, Any]] = []
if episode.source == EpisodeType.message: if episode.source in [EpisodeType.message, EpisodeType.text]:
extracted_node_data = await extract_message_nodes(llm_client, episode, previous_episodes) extracted_node_data = await extract_message_nodes(llm_client, episode, previous_episodes)
elif episode.source == EpisodeType.json: elif episode.source == EpisodeType.json:
extracted_node_data = await extract_json_nodes(llm_client, episode) extracted_node_data = await extract_json_nodes(llm_client, episode)