graphiti-core 0.21.0rc13__py3-none-any.whl → 0.22.0__py3-none-any.whl

graphiti_core/driver/search_interface/search_interface.py

@@ -0,0 +1,89 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class SearchInterface(BaseModel):
+    """
+    This is an interface for implementing custom search logic
+    """
+
+    async def edge_fulltext_search(
+        self,
+        driver: Any,
+        query: str,
+        search_filter: Any,
+        group_ids: list[str] | None = None,
+        limit: int = 100,
+    ) -> list[Any]:
+        raise NotImplementedError
+
+    async def edge_similarity_search(
+        self,
+        driver: Any,
+        search_vector: list[float],
+        source_node_uuid: str | None,
+        target_node_uuid: str | None,
+        search_filter: Any,
+        group_ids: list[str] | None = None,
+        limit: int = 100,
+        min_score: float = 0.7,
+    ) -> list[Any]:
+        raise NotImplementedError
+
+    async def node_fulltext_search(
+        self,
+        driver: Any,
+        query: str,
+        search_filter: Any,
+        group_ids: list[str] | None = None,
+        limit: int = 100,
+    ) -> list[Any]:
+        raise NotImplementedError
+
+    async def node_similarity_search(
+        self,
+        driver: Any,
+        search_vector: list[float],
+        search_filter: Any,
+        group_ids: list[str] | None = None,
+        limit: int = 100,
+        min_score: float = 0.7,
+    ) -> list[Any]:
+        raise NotImplementedError
+
+    async def episode_fulltext_search(
+        self,
+        driver: Any,
+        query: str,
+        search_filter: Any,  # kept for parity even if unused in your impl
+        group_ids: list[str] | None = None,
+        limit: int = 100,
+    ) -> list[Any]:
+        raise NotImplementedError
+
+    # ---------- SEARCH FILTERS (sync) ----------
+    def build_node_search_filters(self, search_filters: Any) -> Any:
+        raise NotImplementedError
+
+    def build_edge_search_filters(self, search_filters: Any) -> Any:
+        raise NotImplementedError
+
+    class Config:
+        arbitrary_types_allowed = True

graphiti_core/edges.py

@@ -25,7 +25,7 @@ from uuid import uuid4
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import ENTITY_EDGE_INDEX_NAME, GraphDriver, GraphProvider
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
 from graphiti_core.helpers import parse_db_date
@@ -53,6 +53,9 @@ class Edge(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.edge_delete(self, driver)
+
         if driver.provider == GraphProvider.KUZU:
             await driver.execute_query(
                 """
@@ -77,17 +80,13 @@ class Edge(BaseModel, ABC):
                 uuid=self.uuid,
             )
 
-            if driver.aoss_client:
-                await driver.aoss_client.delete(
-                    index=ENTITY_EDGE_INDEX_NAME,
-                    id=self.uuid,
-                    params={'routing': self.group_id},
-                )
-
         logger.debug(f'Deleted Edge: {self.uuid}')
 
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.edge_delete_by_uuids(cls, driver, uuids)
+
         if driver.provider == GraphProvider.KUZU:
             await driver.execute_query(
                 """
@@ -115,12 +114,6 @@ class Edge(BaseModel, ABC):
                 uuids=uuids,
             )
 
-            if driver.aoss_client:
-                await driver.aoss_client.delete_by_query(
-                    index=ENTITY_EDGE_INDEX_NAME,
-                    body={'query': {'terms': {'uuid': uuids}}},
-                )
-
         logger.debug(f'Deleted Edges: {uuids}')
 
     def __hash__(self):
@@ -258,6 +251,9 @@ class EntityEdge(Edge):
         return self.fact_embedding
 
     async def load_fact_embedding(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.edge_load_embeddings(self, driver)
+
         query = """
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
             RETURN e.fact_embedding AS fact_embedding
@@ -268,21 +264,6 @@ class EntityEdge(Edge):
                 MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
                 RETURN [x IN split(e.fact_embedding, ",") | toFloat(x)] as fact_embedding
             """
-        elif driver.aoss_client:
-            resp = await driver.aoss_client.search(
-                body={
-                    'query': {'multi_match': {'query': self.uuid, 'fields': ['uuid']}},
-                    'size': 1,
-                },
-                index=ENTITY_EDGE_INDEX_NAME,
-                params={'routing': self.group_id},
-            )
-
-            if resp['hits']['hits']:
-                self.fact_embedding = resp['hits']['hits'][0]['_source']['fact_embedding']
-                return
-            else:
-                raise EdgeNotFoundError(self.uuid)
 
         if driver.provider == GraphProvider.KUZU:
             query = """
@@ -320,15 +301,11 @@ class EntityEdge(Edge):
         if driver.provider == GraphProvider.KUZU:
             edge_data['attributes'] = json.dumps(self.attributes)
             result = await driver.execute_query(
-                get_entity_edge_save_query(driver.provider, has_aoss=bool(driver.aoss_client)),
+                get_entity_edge_save_query(driver.provider),
                 **edge_data,
             )
         else:
             edge_data.update(self.attributes or {})
-
-            if driver.aoss_client:
-                await driver.save_to_aoss(ENTITY_EDGE_INDEX_NAME, [edge_data])  # pyright: ignore reportAttributeAccessIssue
-
             result = await driver.execute_query(
                 get_entity_edge_save_query(driver.provider),
                 edge_data=edge_data,