graphiti-core 0.22.0rc5__py3-none-any.whl → 0.22.1rc1__py3-none-any.whl

graphiti_core/nodes.py

@@ -27,10 +27,6 @@ from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import (
-    COMMUNITY_INDEX_NAME,
-    ENTITY_EDGE_INDEX_NAME,
-    ENTITY_INDEX_NAME,
-    EPISODE_INDEX_NAME,
     GraphDriver,
     GraphProvider,
 )
@@ -99,6 +95,9 @@ class Node(BaseModel, ABC):
     async def save(self, driver: GraphDriver): ...
 
     async def delete(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_delete(self, driver)
+
         match driver.provider:
             case GraphProvider.NEO4J:
                 records, _, _ = await driver.execute_query(
@@ -113,27 +112,6 @@ class Node(BaseModel, ABC):
                     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(
@@ -181,14 +159,18 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_group_id(cls, driver: GraphDriver, group_id: str, batch_size: int = 100):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_delete_by_group_id(
+                cls, driver, group_id, batch_size
+            )
+
         match driver.provider:
             case GraphProvider.NEO4J:
                 async with driver.session() as session:
                     await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community {group_id: $group_id})
-                        CALL {
-                            WITH n
+                        CALL (n) {
                             DETACH DELETE n
                         } IN TRANSACTIONS OF $batch_size ROWS
                         """,
@@ -196,31 +178,6 @@ class Node(BaseModel, ABC):
                         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(
@@ -258,6 +215,11 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str], batch_size: int = 100):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_delete_by_uuids(
+                cls, driver, uuids, group_id=None, batch_size=batch_size
+            )
+
         match driver.provider:
             case GraphProvider.FALKORDB:
                 for label in ['Entity', 'Episodic', 'Community']:
@@ -300,7 +262,7 @@ class Node(BaseModel, ABC):
             case _:  # Neo4J, Neptune
                 async with driver.session() as session:
                     # Collect all edge UUIDs before deleting nodes
-                    result = await session.run(
+                    await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community)
                         WHERE n.uuid IN $uuids
@@ -310,18 +272,12 @@ class Node(BaseModel, ABC):
                         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)
                         WHERE n.uuid IN $uuids
-                        CALL {
-                            WITH n
+                        CALL (n) {
                             DETACH DELETE n
                         } IN TRANSACTIONS OF $batch_size ROWS
                         """,
@@ -329,20 +285,6 @@ 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): ...
 
@@ -363,6 +305,9 @@ class EpisodicNode(Node):
     )
 
     async def save(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.episodic_node_save(self, driver)
+
         episode_args = {
             'uuid': self.uuid,
             'name': self.name,
@@ -375,12 +320,6 @@ 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
         )
@@ -510,26 +449,14 @@ class EntityNode(Node):
         return self.name_embedding
 
     async def load_name_embedding(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_load_embeddings(self, driver)
+
         if driver.provider == GraphProvider.NEPTUNE:
             query: LiteralString = """
                 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 = """
@@ -548,6 +475,9 @@ class EntityNode(Node):
         self.name_embedding = records[0]['name_embedding']
 
     async def save(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_save(self, driver)
+
         entity_data: dict[str, Any] = {
             'uuid': self.uuid,
             'name': self.name,
@@ -568,11 +498,8 @@ class EntityNode(Node):
             entity_data.update(self.attributes or {})
             labels = ':'.join(self.labels + ['Entity'])
 
-            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, bool(driver.aoss_client)),
+                get_entity_node_save_query(driver.provider, labels),
                 entity_data=entity_data,
             )
 

graphiti_core/search/search_filters.py

@@ -249,41 +249,3 @@ 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