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

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(

graphiti_core/search/search_filters.py

@@ -52,6 +52,17 @@ class SearchFilters(BaseModel):
     invalid_at: list[list[DateFilter]] | None = Field(default=None)
     created_at: list[list[DateFilter]] | None = Field(default=None)
     expired_at: list[list[DateFilter]] | None = Field(default=None)
+    edge_uuids: list[str] | None = Field(default=None)
+
+
+def cypher_to_opensearch_operator(op: ComparisonOperator) -> str:
+    mapping = {
+        ComparisonOperator.greater_than: 'gt',
+        ComparisonOperator.less_than: 'lt',
+        ComparisonOperator.greater_than_equal: 'gte',
+        ComparisonOperator.less_than_equal: 'lte',
+    }
+    return mapping.get(op, op.value)
 
 
 def node_search_filter_query_constructor(
@@ -98,6 +109,10 @@ def edge_search_filter_query_constructor(
         filter_queries.append('e.name in $edge_types')
         filter_params['edge_types'] = edge_types
 
+    if filters.edge_uuids is not None:
+        filter_queries.append('e.uuid in $edge_uuids')
+        filter_params['edge_uuids'] = filters.edge_uuids
+
     if filters.node_labels is not None:
         if provider == GraphProvider.KUZU:
             node_label_filter = (
@@ -234,3 +249,41 @@ def edge_search_filter_query_constructor(
         filter_queries.append(expired_at_filter)
 
     return filter_queries, filter_params
+
+
+def build_aoss_node_filters(group_ids: list[str], search_filters: SearchFilters) -> list[dict]:
+    filters = [{'terms': {'group_id': group_ids}}]
+
+    if search_filters.node_labels:
+        filters.append({'terms': {'node_labels': search_filters.node_labels}})
+
+    return filters
+
+
+def build_aoss_edge_filters(group_ids: list[str], search_filters: SearchFilters) -> list[dict]:
+    filters: list[dict] = [{'terms': {'group_id': group_ids}}]
+
+    if search_filters.edge_types:
+        filters.append({'terms': {'edge_types': search_filters.edge_types}})
+
+    if search_filters.edge_uuids:
+        filters.append({'terms': {'uuid': search_filters.edge_uuids}})
+
+    for field in ['valid_at', 'invalid_at', 'created_at', 'expired_at']:
+        ranges = getattr(search_filters, field)
+        if ranges:
+            # OR of ANDs
+            should_clauses = []
+            for and_group in ranges:
+                and_filters = []
+                for df in and_group:  # df is a DateFilter
+                    range_query = {
+                        'range': {
+                            field: {cypher_to_opensearch_operator(df.comparison_operator): df.date}
+                        }
+                    }
+                    and_filters.append(range_query)
+                should_clauses.append({'bool': {'filter': and_filters}})
+            filters.append({'bool': {'should': should_clauses, 'minimum_should_match': 1}})
+
+    return filters