graphiti-core 0.20.4__py3-none-any.whl → 0.21.0__py3-none-any.whl

graphiti_core/llm_client/openai_base_client.py

@@ -25,7 +25,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError, RefusalError
 
@@ -175,6 +175,7 @@ class BaseOpenAIClient(LLMClient):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
     ) -> dict[str, typing.Any]:
         """Generate a response with retry logic and error handling."""
         if max_tokens is None:
@@ -184,7 +185,7 @@ class BaseOpenAIClient(LLMClient):
         last_error = None
 
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         while retry_count <= self.MAX_RETRIES:
             try:

graphiti_core/llm_client/openai_generic_client.py

@@ -25,7 +25,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError, RefusalError
 
@@ -120,6 +120,7 @@ class OpenAIGenericClient(LLMClient):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
     ) -> dict[str, typing.Any]:
         if max_tokens is None:
             max_tokens = self.max_tokens
@@ -136,7 +137,7 @@ class OpenAIGenericClient(LLMClient):
             )
 
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         while retry_count <= self.MAX_RETRIES:
             try:

graphiti_core/models/edges/edge_db_queries.py

@@ -60,7 +60,7 @@ EPISODIC_EDGE_RETURN = """
 """
 
 
-def get_entity_edge_save_query(provider: GraphProvider) -> str:
+def get_entity_edge_save_query(provider: GraphProvider, has_aoss: bool = False) -> str:
     match provider:
         case GraphProvider.FALKORDB:
             return """
@@ -99,17 +99,28 @@ def get_entity_edge_save_query(provider: GraphProvider) -> str:
                 RETURN e.uuid AS uuid
             """
         case _:  # Neo4j
-            return """
-                MATCH (source:Entity {uuid: $edge_data.source_uuid})
-                MATCH (target:Entity {uuid: $edge_data.target_uuid})
-                MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
-                SET e = $edge_data
-                WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)
+            save_embedding_query = (
+                """WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)"""
+                if not has_aoss
+                else ''
+            )
+            return (
+                (
+                    """
+                        MATCH (source:Entity {uuid: $edge_data.source_uuid})
+                        MATCH (target:Entity {uuid: $edge_data.target_uuid})
+                        MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+                        SET e = $edge_data
+                        """
+                    + save_embedding_query
+                )
+                + """
                 RETURN e.uuid AS uuid
-            """
+                """
+            )
 
 
-def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
+def get_entity_edge_save_bulk_query(provider: GraphProvider, has_aoss: bool = False) -> str:
     match provider:
         case GraphProvider.FALKORDB:
             return """
@@ -152,15 +163,24 @@ def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
                 RETURN e.uuid AS uuid
             """
         case _:
-            return """
-                UNWIND $entity_edges AS edge
-                MATCH (source:Entity {uuid: edge.source_node_uuid})
-                MATCH (target:Entity {uuid: edge.target_node_uuid})
-                MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
-                SET e = edge
-                WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)
+            save_embedding_query = (
+                'WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                """
+                    UNWIND $entity_edges AS edge
+                    MATCH (source:Entity {uuid: edge.source_node_uuid})
+                    MATCH (target:Entity {uuid: edge.target_node_uuid})
+                    MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
+                    SET e = edge
+                    """
+                + save_embedding_query
+                + """
                 RETURN edge.uuid AS uuid
             """
+            )
 
 
 def get_entity_edge_return_query(provider: GraphProvider) -> str:

graphiti_core/models/nodes/node_db_queries.py

@@ -126,7 +126,7 @@ EPISODIC_NODE_RETURN_NEPTUNE = """
 """
 
 
-def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
+def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: bool = False) -> str:
     match provider:
         case GraphProvider.FALKORDB:
             return f"""
@@ -161,16 +161,27 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
                 RETURN n.uuid AS uuid
             """
         case _:
-            return f"""
+            save_embedding_query = (
+                'WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                f"""
                 MERGE (n:Entity {{uuid: $entity_data.uuid}})
                 SET n:{labels}
                 SET n = $entity_data
-                WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)
+                """
+                + save_embedding_query
+                + """
                 RETURN n.uuid AS uuid
             """
+            )
 
 
-def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict]) -> str | Any:
+def get_entity_node_save_bulk_query(
+    provider: GraphProvider, nodes: list[dict], has_aoss: bool = False
+) -> str | Any:
     match provider:
         case GraphProvider.FALKORDB:
             queries = []
@@ -222,14 +233,23 @@ def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict])
                 RETURN n.uuid AS uuid
             """
         case _:  # Neo4j
-            return """
-                UNWIND $nodes AS node
-                MERGE (n:Entity {uuid: node.uuid})
-                SET n:$(node.labels)
-                SET n = node
-                WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
+            save_embedding_query = (
+                'WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                """
+                    UNWIND $nodes AS node
+                    MERGE (n:Entity {uuid: node.uuid})
+                    SET n:$(node.labels)
+                    SET n = node
+                    """
+                + save_embedding_query
+                + """
                 RETURN n.uuid AS uuid
             """
+            )
 
 
 def get_entity_node_return_query(provider: GraphProvider) -> str:

graphiti_core/nodes.py

@@ -26,7 +26,14 @@ from uuid import uuid4
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver, GraphProvider
+from graphiti_core.driver.driver import (
+    COMMUNITY_INDEX_NAME,
+    ENTITY_EDGE_INDEX_NAME,
+    ENTITY_INDEX_NAME,
+    EPISODE_INDEX_NAME,
+    GraphDriver,
+    GraphProvider,
+)
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
 from graphiti_core.helpers import parse_db_date
@@ -94,13 +101,39 @@ class Node(BaseModel, ABC):
     async def delete(self, driver: GraphDriver):
         match driver.provider:
             case GraphProvider.NEO4J:
-                await driver.execute_query(
+                records, _, _ = await driver.execute_query(
                     """
-                    MATCH (n:Entity|Episodic|Community {uuid: $uuid})
+                    MATCH (n {uuid: $uuid})
+                    WHERE n:Entity OR n:Episodic OR n:Community
+                    OPTIONAL MATCH (n)-[r]-()
+                    WITH collect(r.uuid) AS edge_uuids, n
                     DETACH DELETE n
+                    RETURN edge_uuids
                     """,
                     uuid=self.uuid,
                 )
+
+                edge_uuids: list[str] = records[0].get('edge_uuids', []) if records else []
+
+                if driver.aoss_client:
+                    # Delete the node from OpenSearch indices
+                    for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME):
+                        await driver.aoss_client.delete(
+                            index=index,
+                            id=self.uuid,
+                            params={'routing': self.group_id},
+                        )
+
+                    # Bulk delete the detached edges
+                    if edge_uuids:
+                        actions = []
+                        for eid in edge_uuids:
+                            actions.append(
+                                {'delete': {'_index': ENTITY_EDGE_INDEX_NAME, '_id': eid}}
+                            )
+
+                        await driver.aoss_client.bulk(body=actions)
+
             case GraphProvider.KUZU:
                 for label in ['Episodic', 'Community']:
                     await driver.execute_query(
@@ -162,6 +195,32 @@ class Node(BaseModel, ABC):
                         group_id=group_id,
                         batch_size=batch_size,
                     )
+
+                if driver.aoss_client:
+                    await driver.aoss_client.delete_by_query(
+                        index=EPISODE_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
+                    await driver.aoss_client.delete_by_query(
+                        index=ENTITY_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
+                    await driver.aoss_client.delete_by_query(
+                        index=COMMUNITY_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
+                    await driver.aoss_client.delete_by_query(
+                        index=ENTITY_EDGE_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
             case GraphProvider.KUZU:
                 for label in ['Episodic', 'Community']:
                     await driver.execute_query(
@@ -240,6 +299,23 @@ class Node(BaseModel, ABC):
                 )
             case _:  # Neo4J, Neptune
                 async with driver.session() as session:
+                    # Collect all edge UUIDs before deleting nodes
+                    result = await session.run(
+                        """
+                        MATCH (n:Entity|Episodic|Community)
+                        WHERE n.uuid IN $uuids
+                        MATCH (n)-[r]-()
+                        RETURN collect(r.uuid) AS edge_uuids
+                        """,
+                        uuids=uuids,
+                    )
+
+                    record = await result.single()
+                    edge_uuids: list[str] = (
+                        record['edge_uuids'] if record and record['edge_uuids'] else []
+                    )
+
+                    # Now delete the nodes in batches
                     await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community)
@@ -253,6 +329,20 @@ class Node(BaseModel, ABC):
                         batch_size=batch_size,
                     )
 
+                if driver.aoss_client:
+                    for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME):
+                        await driver.aoss_client.delete_by_query(
+                            index=index,
+                            body={'query': {'terms': {'uuid': uuids}}},
+                        )
+
+                    if edge_uuids:
+                        actions = [
+                            {'delete': {'_index': ENTITY_EDGE_INDEX_NAME, '_id': eid}}
+                            for eid in edge_uuids
+                        ]
+                        await driver.aoss_client.bulk(body=actions)
+
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
 
@@ -273,20 +363,6 @@ class EpisodicNode(Node):
     )
 
     async def save(self, driver: GraphDriver):
-        if driver.provider == GraphProvider.NEPTUNE:
-            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
-                'episode_content',
-                [
-                    {
-                        'uuid': self.uuid,
-                        'group_id': self.group_id,
-                        'source': self.source.value,
-                        'content': self.content,
-                        'source_description': self.source_description,
-                    }
-                ],
-            )
-
         episode_args = {
             'uuid': self.uuid,
             'name': self.name,
@@ -299,6 +375,12 @@ class EpisodicNode(Node):
             'source': self.source.value,
         }
 
+        if driver.aoss_client:
+            await driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+                'episodes',
+                [episode_args],
+            )
+
         result = await driver.execute_query(
             get_episode_node_save_query(driver.provider), **episode_args
         )
@@ -433,6 +515,22 @@ class EntityNode(Node):
                 MATCH (n:Entity {uuid: $uuid})
                 RETURN [x IN split(n.name_embedding, ",") | toFloat(x)] as name_embedding
             """
+        elif driver.aoss_client:
+            resp = await driver.aoss_client.search(
+                body={
+                    'query': {'multi_match': {'query': self.uuid, 'fields': ['uuid']}},
+                    'size': 1,
+                },
+                index=ENTITY_INDEX_NAME,
+                params={'routing': self.group_id},
+            )
+
+            if resp['hits']['hits']:
+                self.name_embedding = resp['hits']['hits'][0]['_source']['name_embedding']
+                return
+            else:
+                raise NodeNotFoundError(self.uuid)
+
         else:
             query: LiteralString = """
                 MATCH (n:Entity {uuid: $uuid})
@@ -470,11 +568,11 @@ class EntityNode(Node):
             entity_data.update(self.attributes or {})
             labels = ':'.join(self.labels + ['Entity'])
 
-            if driver.provider == GraphProvider.NEPTUNE:
-                driver.save_to_aoss('node_name_and_summary', [entity_data])  # pyright: ignore reportAttributeAccessIssue
+            if driver.aoss_client:
+                await driver.save_to_aoss(ENTITY_INDEX_NAME, [entity_data])  # pyright: ignore reportAttributeAccessIssue
 
             result = await driver.execute_query(
-                get_entity_node_save_query(driver.provider, labels),
+                get_entity_node_save_query(driver.provider, labels, bool(driver.aoss_client)),
                 entity_data=entity_data,
             )
 
@@ -569,8 +667,8 @@ class CommunityNode(Node):
 
     async def save(self, driver: GraphDriver):
         if driver.provider == GraphProvider.NEPTUNE:
-            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
-                'community_name',
+            await driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+                'communities',
                 [{'name': self.name, 'uuid': self.uuid, 'group_id': self.group_id}],
             )
         result = await driver.execute_query(
@@ -770,9 +868,12 @@ def get_community_node_from_record(record: Any) -> CommunityNode:
 
 
 async def create_entity_node_embeddings(embedder: EmbedderClient, nodes: list[EntityNode]):
-    if not nodes:  # Handle empty list case
+    # filter out falsey values from nodes
+    filtered_nodes = [node for node in nodes if node.name]
+
+    if not filtered_nodes:
         return
 
-    name_embeddings = await embedder.create_batch([node.name for node in nodes])
-    for node, name_embedding in zip(nodes, name_embeddings, strict=True):
+    name_embeddings = await embedder.create_batch([node.name for node in filtered_nodes])
+    for node, name_embedding in zip(filtered_nodes, name_embeddings, strict=True):
         node.name_embedding = name_embedding

graphiti_core/prompts/dedupe_edges.py

@@ -25,11 +25,11 @@ from .prompt_helpers import to_prompt_json
 class EdgeDuplicate(BaseModel):
     duplicate_facts: list[int] = Field(
         ...,
-        description='List of ids of any duplicate facts. If no duplicate facts are found, default to empty list.',
+        description='List of idx values of any duplicate facts. If no duplicate facts are found, default to empty list.',
     )
     contradicted_facts: list[int] = Field(
         ...,
-        description='List of ids of facts that should be invalidated. If no facts should be invalidated, the list should be empty.',
+        description='List of idx values of facts that should be invalidated. If no facts should be invalidated, the list should be empty.',
     )
     fact_type: str = Field(..., description='One of the provided fact types or DEFAULT')
 
@@ -67,11 +67,11 @@ def edge(context: dict[str, Any]) -> list[Message]:
         Given the following context, determine whether the New Edge represents any of the edges in the list of Existing Edges.
 
         <EXISTING EDGES>
-        {to_prompt_json(context['related_edges'], ensure_ascii=context.get('ensure_ascii', False), indent=2)}
+        {to_prompt_json(context['related_edges'], indent=2)}
         </EXISTING EDGES>
 
         <NEW EDGE>
-        {to_prompt_json(context['extracted_edges'], ensure_ascii=context.get('ensure_ascii', False), indent=2)}
+        {to_prompt_json(context['extracted_edges'], indent=2)}
         </NEW EDGE>
         
         Task:
@@ -98,7 +98,7 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
         Given the following context, find all of the duplicates in a list of facts:
 
         Facts:
-        {to_prompt_json(context['edges'], ensure_ascii=context.get('ensure_ascii', False), indent=2)}
+        {to_prompt_json(context['edges'], indent=2)}
 
         Task:
         If any facts in Facts is a duplicate of another fact, return a new fact with one of their uuid's.
@@ -124,37 +124,48 @@ def resolve_edge(context: dict[str, Any]) -> list[Message]:
         Message(
             role='user',
             content=f"""
-        <NEW FACT>
-        {context['new_edge']}
-        </NEW FACT>
-        
+        Task:
+        You will receive TWO separate lists of facts. Each list uses 'idx' as its index field, starting from 0.
+
+        1. DUPLICATE DETECTION:
+           - If the NEW FACT represents identical factual information as any fact in EXISTING FACTS, return those idx values in duplicate_facts.
+           - Facts with similar information that contain key differences should NOT be marked as duplicates.
+           - Return idx values from EXISTING FACTS.
+           - If no duplicates, return an empty list for duplicate_facts.
+
+        2. FACT TYPE CLASSIFICATION:
+           - Given the predefined FACT TYPES, determine if the NEW FACT should be classified as one of these types.
+           - Return the fact type as fact_type or DEFAULT if NEW FACT is not one of the FACT TYPES.
+
+        3. CONTRADICTION DETECTION:
+           - Based on FACT INVALIDATION CANDIDATES and NEW FACT, determine which facts the new fact contradicts.
+           - Return idx values from FACT INVALIDATION CANDIDATES.
+           - If no contradictions, return an empty list for contradicted_facts.
+
+        IMPORTANT:
+        - duplicate_facts: Use ONLY 'idx' values from EXISTING FACTS
+        - contradicted_facts: Use ONLY 'idx' values from FACT INVALIDATION CANDIDATES
+        - These are two separate lists with independent idx ranges starting from 0
+
+        Guidelines:
+        1. Some facts may be very similar but will have key differences, particularly around numeric values in the facts.
+            Do not mark these facts as duplicates.
+
+        <FACT TYPES>
+        {context['edge_types']}
+        </FACT TYPES>
+
         <EXISTING FACTS>
         {context['existing_edges']}
         </EXISTING FACTS>
+
         <FACT INVALIDATION CANDIDATES>
         {context['edge_invalidation_candidates']}
         </FACT INVALIDATION CANDIDATES>
-        
-        <FACT TYPES>
-        {context['edge_types']}
-        </FACT TYPES>
-        
 
-        Task:
-        If the NEW FACT represents identical factual information of one or more in EXISTING FACTS, return the idx of the duplicate facts.
-        Facts with similar information that contain key differences should not be marked as duplicates.
-        If the NEW FACT is not a duplicate of any of the EXISTING FACTS, return an empty list.
-        
-        Given the predefined FACT TYPES, determine if the NEW FACT should be classified as one of these types.
-        Return the fact type as fact_type or DEFAULT if NEW FACT is not one of the FACT TYPES.
-        
-        Based on the provided FACT INVALIDATION CANDIDATES and NEW FACT, determine which existing facts the new fact contradicts.
-        Return a list containing all idx's of the facts that are contradicted by the NEW FACT.
-        If there are no contradicted facts, return an empty list.
-
-        Guidelines:
-        1. Some facts may be very similar but will have key differences, particularly around numeric values in the facts.
-            Do not mark these facts as duplicates.
+        <NEW FACT>
+        {context['new_edge']}
+        </NEW FACT>
         """,
         ),
     ]