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

graphiti_core/driver/driver.py

@@ -24,6 +24,9 @@ from typing import Any
 
 from dotenv import load_dotenv
 
+from graphiti_core.driver.graph_operations.graph_operations import GraphOperationsInterface
+from graphiti_core.driver.search_interface.search_interface import SearchInterface
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_SIZE = 10
@@ -73,7 +76,8 @@ class GraphDriver(ABC):
         ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
     )
     _database: str
-    aoss_client: Any  # type: ignore
+    search_interface: SearchInterface | None = None
+    graph_operations_interface: GraphOperationsInterface | None = None
 
     @abstractmethod
     def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
@@ -109,9 +113,3 @@ class GraphDriver(ABC):
         Only implemented by providers that need custom fulltext query building.
         """
         raise NotImplementedError(f'build_fulltext_query not implemented for {self.provider}')
-
-    async def save_to_aoss(self, name: str, data: list[dict]) -> int:
-        return 0
-
-    async def clear_aoss_indices(self):
-        return 1

graphiti_core/driver/falkordb_driver.py

@@ -14,6 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import asyncio
+import datetime
 import logging
 from typing import TYPE_CHECKING, Any
 
@@ -191,9 +193,36 @@ class FalkorDriver(GraphDriver):
             await self.client.connection.close()
 
     async def delete_all_indexes(self) -> None:
-        await self.execute_query(
-            'CALL db.indexes() YIELD name DROP INDEX name',
-        )
+        result = await self.execute_query('CALL db.indexes()')
+        if not result:
+            return
+
+        records, _, _ = result
+        drop_tasks = []
+
+        for record in records:
+            label = record['label']
+            entity_type = record['entitytype']
+
+            for field_name, index_type in record['types'].items():
+                if 'RANGE' in index_type:
+                    drop_tasks.append(self.execute_query(f'DROP INDEX ON :{label}({field_name})'))
+                elif 'FULLTEXT' in index_type:
+                    if entity_type == 'NODE':
+                        drop_tasks.append(
+                            self.execute_query(
+                                f'DROP FULLTEXT INDEX FOR (n:{label}) ON (n.{field_name})'
+                            )
+                        )
+                    elif entity_type == 'RELATIONSHIP':
+                        drop_tasks.append(
+                            self.execute_query(
+                                f'DROP FULLTEXT INDEX FOR ()-[e:{label}]-() ON (e.{field_name})'
+                            )
+                        )
+
+        if drop_tasks:
+            await asyncio.gather(*drop_tasks)
 
     def clone(self, database: str) -> 'GraphDriver':
         """
@@ -204,6 +233,28 @@ class FalkorDriver(GraphDriver):
 
         return cloned
 
+    async def health_check(self) -> None:
+        """Check FalkorDB connectivity by running a simple query."""
+        try:
+            await self.execute_query('MATCH (n) RETURN 1 LIMIT 1')
+            return None
+        except Exception as e:
+            print(f'FalkorDB health check failed: {e}')
+            raise
+
+    @staticmethod
+    def convert_datetimes_to_strings(obj):
+        if isinstance(obj, dict):
+            return {k: FalkorDriver.convert_datetimes_to_strings(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [FalkorDriver.convert_datetimes_to_strings(item) for item in obj]
+        elif isinstance(obj, tuple):
+            return tuple(FalkorDriver.convert_datetimes_to_strings(item) for item in obj)
+        elif isinstance(obj, datetime):
+            return obj.isoformat()
+        else:
+            return obj
+
     def sanitize(self, query: str) -> str:
         """
         Replace FalkorDB special characters with whitespace.

graphiti_core/driver/graph_operations/graph_operations.py

@@ -0,0 +1,195 @@
+"""
+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 GraphOperationsInterface(BaseModel):
+    """
+    Interface for updating graph mutation behavior.
+    """
+
+    # -----------------
+    # Node: Save/Delete
+    # -----------------
+
+    async def node_save(self, node: Any, driver: Any) -> None:
+        """Persist (create or update) a single node."""
+        raise NotImplementedError
+
+    async def node_delete(self, node: Any, driver: Any) -> None:
+        raise NotImplementedError
+
+    async def node_save_bulk(
+        self,
+        _cls: Any,  # kept for parity; callers won't pass it
+        driver: Any,
+        transaction: Any,
+        nodes: list[Any],
+        batch_size: int = 100,
+    ) -> None:
+        """Persist (create or update) many nodes in batches."""
+        raise NotImplementedError
+
+    async def node_delete_by_group_id(
+        self,
+        _cls: Any,
+        driver: Any,
+        group_id: str,
+        batch_size: int = 100,
+    ) -> None:
+        raise NotImplementedError
+
+    async def node_delete_by_uuids(
+        self,
+        _cls: Any,
+        driver: Any,
+        uuids: list[str],
+        group_id: str | None = None,
+        batch_size: int = 100,
+    ) -> None:
+        raise NotImplementedError
+
+    # --------------------------
+    # Node: Embeddings (load)
+    # --------------------------
+
+    async def node_load_embeddings(self, node: Any, driver: Any) -> None:
+        """
+        Load embedding vectors for a single node into the instance (e.g., set node.embedding or similar).
+        """
+        raise NotImplementedError
+
+    async def node_load_embeddings_bulk(
+        self,
+        _cls: Any,
+        driver: Any,
+        transaction: Any,
+        nodes: list[Any],
+        batch_size: int = 100,
+    ) -> None:
+        """
+        Load embedding vectors for many nodes in batches. Mutates the provided node instances.
+        """
+        raise NotImplementedError
+
+    # --------------------------
+    # EpisodicNode: Save/Delete
+    # --------------------------
+
+    async def episodic_node_save(self, node: Any, driver: Any) -> None:
+        """Persist (create or update) a single episodic node."""
+        raise NotImplementedError
+
+    async def episodic_node_delete(self, node: Any, driver: Any) -> None:
+        raise NotImplementedError
+
+    async def episodic_node_save_bulk(
+        self,
+        _cls: Any,
+        driver: Any,
+        transaction: Any,
+        nodes: list[Any],
+        batch_size: int = 100,
+    ) -> None:
+        """Persist (create or update) many episodic nodes in batches."""
+        raise NotImplementedError
+
+    async def episodic_edge_save_bulk(
+        self,
+        _cls: Any,
+        driver: Any,
+        transaction: Any,
+        episodic_edges: list[Any],
+        batch_size: int = 100,
+    ) -> None:
+        """Persist (create or update) many episodic edges in batches."""
+        raise NotImplementedError
+
+    async def episodic_node_delete_by_group_id(
+        self,
+        _cls: Any,
+        driver: Any,
+        group_id: str,
+        batch_size: int = 100,
+    ) -> None:
+        raise NotImplementedError
+
+    async def episodic_node_delete_by_uuids(
+        self,
+        _cls: Any,
+        driver: Any,
+        uuids: list[str],
+        group_id: str | None = None,
+        batch_size: int = 100,
+    ) -> None:
+        raise NotImplementedError
+
+    # -----------------
+    # Edge: Save/Delete
+    # -----------------
+
+    async def edge_save(self, edge: Any, driver: Any) -> None:
+        """Persist (create or update) a single edge."""
+        raise NotImplementedError
+
+    async def edge_delete(self, edge: Any, driver: Any) -> None:
+        raise NotImplementedError
+
+    async def edge_save_bulk(
+        self,
+        _cls: Any,
+        driver: Any,
+        transaction: Any,
+        edges: list[Any],
+        batch_size: int = 100,
+    ) -> None:
+        """Persist (create or update) many edges in batches."""
+        raise NotImplementedError
+
+    async def edge_delete_by_uuids(
+        self,
+        _cls: Any,
+        driver: Any,
+        uuids: list[str],
+        group_id: str | None = None,
+    ) -> None:
+        raise NotImplementedError
+
+    # -----------------
+    # Edge: Embeddings (load)
+    # -----------------
+
+    async def edge_load_embeddings(self, edge: Any, driver: Any) -> None:
+        """
+        Load embedding vectors for a single edge into the instance (e.g., set edge.embedding or similar).
+        """
+        raise NotImplementedError
+
+    async def edge_load_embeddings_bulk(
+        self,
+        _cls: Any,
+        driver: Any,
+        transaction: Any,
+        edges: list[Any],
+        batch_size: int = 100,
+    ) -> None:
+        """
+        Load embedding vectors for many edges in batches. Mutates the provided edge instances.
+        """
+        raise NotImplementedError

graphiti_core/driver/neo4j_driver.py

@@ -72,3 +72,12 @@ class Neo4jDriver(GraphDriver):
         return self.client.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
         )
+
+    async def health_check(self) -> None:
+        """Check Neo4j connectivity by running the driver's verify_connectivity method."""
+        try:
+            await self.client.verify_connectivity()
+            return None
+        except Exception as e:
+            print(f'Neo4j health check failed: {e}')
+            raise

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,

graphiti_core/models/edges/edge_db_queries.py

@@ -68,6 +68,7 @@ def get_entity_edge_save_query(provider: GraphProvider, has_aoss: bool = False)
                 MATCH (target:Entity {uuid: $edge_data.target_uuid})
                 MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
                 SET e = $edge_data
+                SET e.fact_embedding = vecf32($edge_data.fact_embedding)
                 RETURN e.uuid AS uuid
             """
         case GraphProvider.NEPTUNE:

graphiti_core/models/nodes/node_db_queries.py

@@ -133,6 +133,7 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: b
                 MERGE (n:Entity {{uuid: $entity_data.uuid}})
                 SET n:{labels}
                 SET n = $entity_data
+                SET n.name_embedding = vecf32($entity_data.name_embedding)
                 RETURN n.uuid AS uuid
             """
         case GraphProvider.KUZU: