graphiti-core 0.17.4__py3-none-any.whl → 0.25.3__py3-none-any.whl

graphiti_core/driver/graph_operations/graph_operations.py

@@ -0,0 +1,191 @@
+"""
+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,
+        driver: Any,
+        nodes: list[Any],
+        batch_size: int = 100,
+    ) -> dict[str, list[float]]:
+        """
+        Load embedding vectors for many nodes in batches.
+        """
+        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,
+        driver: Any,
+        edges: list[Any],
+        batch_size: int = 100,
+    ) -> dict[str, list[float]]:
+        """
+        Load embedding vectors for many edges in batches
+        """
+        raise NotImplementedError

graphiti_core/driver/kuzu_driver.py

@@ -0,0 +1,182 @@
+"""
+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.
+"""
+
+import logging
+from typing import Any
+
+import kuzu
+
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+
+logger = logging.getLogger(__name__)
+
+# Kuzu requires an explicit schema.
+# As Kuzu currently does not support creating full text indexes on edge properties,
+# we work around this by representing (n:Entity)-[:RELATES_TO]->(m:Entity) as
+# (n)-[:RELATES_TO]->(e:RelatesToNode_)-[:RELATES_TO]->(m).
+SCHEMA_QUERIES = """
+    CREATE NODE TABLE IF NOT EXISTS Episodic (
+        uuid STRING PRIMARY KEY,
+        name STRING,
+        group_id STRING,
+        created_at TIMESTAMP,
+        source STRING,
+        source_description STRING,
+        content STRING,
+        valid_at TIMESTAMP,
+        entity_edges STRING[]
+    );
+    CREATE NODE TABLE IF NOT EXISTS Entity (
+        uuid STRING PRIMARY KEY,
+        name STRING,
+        group_id STRING,
+        labels STRING[],
+        created_at TIMESTAMP,
+        name_embedding FLOAT[],
+        summary STRING,
+        attributes STRING
+    );
+    CREATE NODE TABLE IF NOT EXISTS Community (
+        uuid STRING PRIMARY KEY,
+        name STRING,
+        group_id STRING,
+        created_at TIMESTAMP,
+        name_embedding FLOAT[],
+        summary STRING
+    );
+    CREATE NODE TABLE IF NOT EXISTS RelatesToNode_ (
+        uuid STRING PRIMARY KEY,
+        group_id STRING,
+        created_at TIMESTAMP,
+        name STRING,
+        fact STRING,
+        fact_embedding FLOAT[],
+        episodes STRING[],
+        expired_at TIMESTAMP,
+        valid_at TIMESTAMP,
+        invalid_at TIMESTAMP,
+        attributes STRING
+    );
+    CREATE REL TABLE IF NOT EXISTS RELATES_TO(
+        FROM Entity TO RelatesToNode_,
+        FROM RelatesToNode_ TO Entity
+    );
+    CREATE REL TABLE IF NOT EXISTS MENTIONS(
+        FROM Episodic TO Entity,
+        uuid STRING PRIMARY KEY,
+        group_id STRING,
+        created_at TIMESTAMP
+    );
+    CREATE REL TABLE IF NOT EXISTS HAS_MEMBER(
+        FROM Community TO Entity,
+        FROM Community TO Community,
+        uuid STRING,
+        group_id STRING,
+        created_at TIMESTAMP
+    );
+"""
+
+
+class KuzuDriver(GraphDriver):
+    provider: GraphProvider = GraphProvider.KUZU
+    aoss_client: None = None
+
+    def __init__(
+        self,
+        db: str = ':memory:',
+        max_concurrent_queries: int = 1,
+    ):
+        super().__init__()
+        self.db = kuzu.Database(db)
+
+        self.setup_schema()
+
+        self.client = kuzu.AsyncConnection(self.db, max_concurrent_queries=max_concurrent_queries)
+
+    async def execute_query(
+        self, cypher_query_: str, **kwargs: Any
+    ) -> tuple[list[dict[str, Any]] | list[list[dict[str, Any]]], None, None]:
+        params = {k: v for k, v in kwargs.items() if v is not None}
+        # Kuzu does not support these parameters.
+        params.pop('database_', None)
+        params.pop('routing_', None)
+
+        try:
+            results = await self.client.execute(cypher_query_, parameters=params)
+        except Exception as e:
+            params = {k: (v[:5] if isinstance(v, list) else v) for k, v in params.items()}
+            logger.error(f'Error executing Kuzu query: {e}\n{cypher_query_}\n{params}')
+            raise
+
+        if not results:
+            return [], None, None
+
+        if isinstance(results, list):
+            dict_results = [list(result.rows_as_dict()) for result in results]
+        else:
+            dict_results = list(results.rows_as_dict())
+        return dict_results, None, None  # type: ignore
+
+    def session(self, _database: str | None = None) -> GraphDriverSession:
+        return KuzuDriverSession(self)
+
+    async def close(self):
+        # Do not explicitly close the connection, instead rely on GC.
+        pass
+
+    def delete_all_indexes(self, database_: str):
+        pass
+
+    async def build_indices_and_constraints(self, delete_existing: bool = False):
+        # Kuzu doesn't support dynamic index creation like Neo4j or FalkorDB
+        # Schema and indices are created during setup_schema()
+        # This method is required by the abstract base class but is a no-op for Kuzu
+        pass
+
+    def setup_schema(self):
+        conn = kuzu.Connection(self.db)
+        conn.execute(SCHEMA_QUERIES)
+        conn.close()
+
+
+class KuzuDriverSession(GraphDriverSession):
+    provider = GraphProvider.KUZU
+
+    def __init__(self, driver: KuzuDriver):
+        self.driver = driver
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Kuzu, but method must exist.
+        pass
+
+    async def close(self):
+        # Do not close the session here, as we're reusing the driver connection.
+        pass
+
+    async def execute_write(self, func, *args, **kwargs):
+        # Directly await the provided async function with `self` as the transaction/session
+        return await func(self, *args, **kwargs)
+
+    async def run(self, query: str | list, **kwargs: Any) -> Any:
+        if isinstance(query, list):
+            for cypher, params in query:
+                await self.driver.execute_query(cypher, **params)
+        else:
+            await self.driver.execute_query(query, **kwargs)
+        return None

graphiti_core/driver/neo4j_driver.py

@@ -19,17 +19,27 @@ from collections.abc import Coroutine
 from typing import Any
 
 from neo4j import AsyncGraphDatabase, EagerResult
+from neo4j.exceptions import ClientError
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+from graphiti_core.graph_queries import get_fulltext_indices, get_range_indices
+from graphiti_core.helpers import semaphore_gather
 
 logger = logging.getLogger(__name__)
 
 
 class Neo4jDriver(GraphDriver):
-    provider: str = 'neo4j'
-
-    def __init__(self, uri: str, user: str | None, password: str | None, database: str = 'neo4j'):
+    provider = GraphProvider.NEO4J
+    default_group_id: str = ''
+
+    def __init__(
+        self,
+        uri: str,
+        user: str | None,
+        password: str | None,
+        database: str = 'neo4j',
+    ):
         super().__init__()
         self.client = AsyncGraphDatabase.driver(
             uri=uri,
@@ -37,6 +47,20 @@ class Neo4jDriver(GraphDriver):
         )
         self._database = database
 
+        # Schedule the indices and constraints to be built
+        import asyncio
+
+        try:
+            # Try to get the current event loop
+            loop = asyncio.get_running_loop()
+            # Schedule the build_indices_and_constraints to run
+            loop.create_task(self.build_indices_and_constraints())
+        except RuntimeError:
+            # No event loop running, this will be handled later
+            pass
+
+        self.aoss_client = None
+
     async def execute_query(self, cypher_query_: LiteralString, **kwargs: Any) -> EagerResult:
         # Check if database_ is provided in kwargs.
         # If not populated, set the value to retain backwards compatibility
@@ -45,7 +69,11 @@ class Neo4jDriver(GraphDriver):
             params = {}
         params.setdefault('database_', self._database)
 
-        result = await self.client.execute_query(cypher_query_, parameters_=params, **kwargs)
+        try:
+            result = await self.client.execute_query(cypher_query_, parameters_=params, **kwargs)
+        except Exception as e:
+            logger.error(f'Error executing Neo4j query: {e}\n{cypher_query_}\n{params}')
+            raise
 
         return result
 
@@ -56,9 +84,43 @@ class Neo4jDriver(GraphDriver):
     async def close(self) -> None:
         return await self.client.close()
 
-    def delete_all_indexes(self, database_: str | None = None) -> Coroutine[Any, Any, EagerResult]:
-        database = database_ or self._database
+    def delete_all_indexes(self) -> Coroutine:
         return self.client.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
-            database_=database,
         )
+
+    async def _execute_index_query(self, query: LiteralString) -> EagerResult | None:
+        """Execute an index creation query, ignoring 'index already exists' errors.
+
+        Neo4j can raise EquivalentSchemaRuleAlreadyExists when concurrent CREATE INDEX
+        IF NOT EXISTS queries race, even though the index exists. This is safe to ignore.
+        """
+        try:
+            return await self.execute_query(query)
+        except ClientError as e:
+            # Ignore "equivalent index already exists" error (race condition with IF NOT EXISTS)
+            if 'EquivalentSchemaRuleAlreadyExists' in str(e):
+                logger.debug(f'Index already exists (concurrent creation): {query[:50]}...')
+                return None
+            raise
+
+    async def build_indices_and_constraints(self, delete_existing: bool = False):
+        if delete_existing:
+            await self.delete_all_indexes()
+
+        range_indices: list[LiteralString] = get_range_indices(self.provider)
+
+        fulltext_indices: list[LiteralString] = get_fulltext_indices(self.provider)
+
+        index_queries: list[LiteralString] = range_indices + fulltext_indices
+
+        await semaphore_gather(*[self._execute_index_query(query) for query in index_queries])
+
+    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