graphiti-core 0.7.7__tar.gz → 0.7.9__tar.gz

{graphiti_core-0.7.7 → graphiti_core-0.7.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.7.7
+Version: 0.7.9
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.7.7 → graphiti_core-0.7.9}/graphiti_core/prompts/extract_nodes.py

@@ -30,14 +30,17 @@ class MissedEntities(BaseModel):
     missed_entities: list[str] = Field(..., description="Names of entities that weren't extracted")
 
 
-class EntityClassification(BaseModel):
-    entities: list[str] = Field(
-        ...,
-        description='List of entities',
+class EntityClassificationTriple(BaseModel):
+    uuid: str = Field(description='UUID of the entity')
+    name: str = Field(description='Name of the entity')
+    entity_type: str | None = Field(
+        default=None, description='Type of the entity. Must be one of the provided types or None'
     )
-    entity_classifications: list[str | None] = Field(
-        ...,
-        description='List of entities classifications. The index of the classification should match the index of the entity it corresponds to.',
+
+
+class EntityClassification(BaseModel):
+    entity_classifications: list[EntityClassificationTriple] = Field(
+        ..., description='List of entities classification triples.'
     )
 
 
@@ -180,7 +183,7 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
     {context['entity_types']}
     </ENTITY TYPES>
     
-    Given the above conversation, extracted entities, and provided entity types, classify the extracted entities.
+    Given the above conversation, extracted entities, and provided entity types and their descriptions, classify the extracted entities.
     
     Guidelines:
     1. Each entity must have exactly one type

{graphiti_core-0.7.7 → graphiti_core-0.7.9}/graphiti_core/prompts/summarize_nodes.py

@@ -85,6 +85,8 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         provided ENTITY. Summaries must be under 500 words.
         
         In addition, extract any values for the provided entity properties based on their descriptions.
+        If the value of the entity property cannot be found in the current context, set the value of the property to None.
+        Do not hallucinate entity property values if they cannot be found in the current context.
         
         <ENTITY>
         {context['node_name']}

{graphiti_core-0.7.7 → graphiti_core-0.7.9}/graphiti_core/utils/maintenance/node_operations.py

@@ -25,7 +25,11 @@ from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.prompts import prompt_library
 from graphiti_core.prompts.dedupe_nodes import NodeDuplicate
-from graphiti_core.prompts.extract_nodes import EntityClassification, ExtractedNodes, MissedEntities
+from graphiti_core.prompts.extract_nodes import (
+    EntityClassification,
+    ExtractedNodes,
+    MissedEntities,
+)
 from graphiti_core.prompts.summarize_nodes import Summary
 from graphiti_core.utils.datetime_utils import utc_now
 
@@ -152,7 +156,12 @@ async def extract_nodes(
         'episode_content': episode.content,
         'previous_episodes': [ep.content for ep in previous_episodes],
         'extracted_entities': extracted_node_names,
-        'entity_types': entity_types.keys() if entity_types is not None else [],
+        'entity_types': {
+            type_name: values.model_json_schema().get('description')
+            for type_name, values in entity_types.items()
+        }
+        if entity_types is not None
+        else {},
     }
 
     node_classifications: dict[str, str | None] = {}
@@ -163,9 +172,13 @@ async def extract_nodes(
                 prompt_library.extract_nodes.classify_nodes(node_classification_context),
                 response_model=EntityClassification,
             )
-            entities = llm_response.get('entities', [])
             entity_classifications = llm_response.get('entity_classifications', [])
-            node_classifications.update(dict(zip(entities, entity_classifications)))
+            node_classifications.update(
+                {
+                    entity_classification.get('name'): entity_classification.get('entity_type')
+                    for entity_classification in entity_classifications
+                }
+            )
         # catch classification errors and continue if we can't classify
         except Exception as e:
             logger.exception(e)

{graphiti_core-0.7.7 → graphiti_core-0.7.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.7.7"
+version = "0.7.9"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <paul@getzep.com>",