graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

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

@@ -0,0 +1,117 @@
+"""
+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 collections.abc import Coroutine
+from typing import Any
+
+from neo4j import AsyncGraphDatabase, EagerResult
+from typing_extensions import LiteralString
+
+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 = 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,
+            auth=(user or '', password or ''),
+        )
+        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
+        params = kwargs.pop('params', None)
+        if params is None:
+            params = {}
+        params.setdefault('database_', self._database)
+
+        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
+
+    def session(self, database: str | None = None) -> GraphDriverSession:
+        _database = database or self._database
+        return self.client.session(database=_database)  # type: ignore
+
+    async def close(self) -> None:
+        return await self.client.close()
+
+    def delete_all_indexes(self) -> Coroutine:
+        return self.client.execute_query(
+            'CALL db.indexes() YIELD name DROP INDEX name',
+        )
+
+    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_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

graphiti_core/driver/neptune_driver.py

@@ -0,0 +1,305 @@
+"""
+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 asyncio
+import datetime
+import logging
+from collections.abc import Coroutine
+from typing import Any
+
+import boto3
+from langchain_aws.graphs import NeptuneAnalyticsGraph, NeptuneGraph
+from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers
+
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+
+logger = logging.getLogger(__name__)
+DEFAULT_SIZE = 10
+
+aoss_indices = [
+    {
+        'index_name': 'node_name_and_summary',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'summary': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {'multi_match': {'query': '', 'fields': ['name', 'summary', 'group_id']}},
+            'size': DEFAULT_SIZE,
+        },
+    },
+    {
+        'index_name': 'community_name',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {'multi_match': {'query': '', 'fields': ['name', 'group_id']}},
+            'size': DEFAULT_SIZE,
+        },
+    },
+    {
+        'index_name': 'episode_content',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'content': {'type': 'text'},
+                    'source': {'type': 'text'},
+                    'source_description': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {
+                'multi_match': {
+                    'query': '',
+                    'fields': ['content', 'source', 'source_description', 'group_id'],
+                }
+            },
+            'size': DEFAULT_SIZE,
+        },
+    },
+    {
+        'index_name': 'edge_name_and_fact',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'fact': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+        'query': {
+            'query': {'multi_match': {'query': '', 'fields': ['name', 'fact', 'group_id']}},
+            'size': DEFAULT_SIZE,
+        },
+    },
+]
+
+
+class NeptuneDriver(GraphDriver):
+    provider: GraphProvider = GraphProvider.NEPTUNE
+
+    def __init__(self, host: str, aoss_host: str, port: int = 8182, aoss_port: int = 443):
+        """This initializes a NeptuneDriver for use with Neptune as a backend
+
+        Args:
+            host (str): The Neptune Database or Neptune Analytics host
+            aoss_host (str): The OpenSearch host value
+            port (int, optional): The Neptune Database port, ignored for Neptune Analytics. Defaults to 8182.
+            aoss_port (int, optional): The OpenSearch port. Defaults to 443.
+        """
+        if not host:
+            raise ValueError('You must provide an endpoint to create a NeptuneDriver')
+
+        if host.startswith('neptune-db://'):
+            # This is a Neptune Database Cluster
+            endpoint = host.replace('neptune-db://', '')
+            self.client = NeptuneGraph(endpoint, port)
+            logger.debug('Creating Neptune Database session for %s', host)
+        elif host.startswith('neptune-graph://'):
+            # This is a Neptune Analytics Graph
+            graphId = host.replace('neptune-graph://', '')
+            self.client = NeptuneAnalyticsGraph(graphId)
+            logger.debug('Creating Neptune Graph session for %s', host)
+        else:
+            raise ValueError(
+                'You must provide an endpoint to create a NeptuneDriver as either neptune-db://<endpoint> or neptune-graph://<graphid>'
+            )
+
+        if not aoss_host:
+            raise ValueError('You must provide an AOSS endpoint to create an OpenSearch driver.')
+
+        session = boto3.Session()
+        self.aoss_client = OpenSearch(
+            hosts=[{'host': aoss_host, 'port': aoss_port}],
+            http_auth=Urllib3AWSV4SignerAuth(
+                session.get_credentials(), session.region_name, 'aoss'
+            ),
+            use_ssl=True,
+            verify_certs=True,
+            connection_class=Urllib3HttpConnection,
+            pool_maxsize=20,
+        )
+
+    def _sanitize_parameters(self, query, params: dict):
+        if isinstance(query, list):
+            queries = []
+            for q in query:
+                queries.append(self._sanitize_parameters(q, params))
+            return queries
+        else:
+            for k, v in params.items():
+                if isinstance(v, datetime.datetime):
+                    params[k] = v.isoformat()
+                elif isinstance(v, list):
+                    # Handle lists that might contain datetime objects
+                    for i, item in enumerate(v):
+                        if isinstance(item, datetime.datetime):
+                            v[i] = item.isoformat()
+                            query = str(query).replace(f'${k}', f'datetime(${k})')
+                        if isinstance(item, dict):
+                            query = self._sanitize_parameters(query, v[i])
+
+                    # If the list contains datetime objects, we need to wrap each element with datetime()
+                    if any(isinstance(item, str) and 'T' in item for item in v):
+                        # Create a new list expression with datetime() wrapped around each element
+                        datetime_list = (
+                            '['
+                            + ', '.join(
+                                f'datetime("{item}")'
+                                if isinstance(item, str) and 'T' in item
+                                else repr(item)
+                                for item in v
+                            )
+                            + ']'
+                        )
+                        query = str(query).replace(f'${k}', datetime_list)
+                elif isinstance(v, dict):
+                    query = self._sanitize_parameters(query, v)
+            return query
+
+    async def execute_query(
+        self, cypher_query_, **kwargs: Any
+    ) -> tuple[dict[str, Any], None, None]:
+        params = dict(kwargs)
+        if isinstance(cypher_query_, list):
+            for q in cypher_query_:
+                result, _, _ = self._run_query(q[0], q[1])
+            return result, None, None
+        else:
+            return self._run_query(cypher_query_, params)
+
+    def _run_query(self, cypher_query_, params):
+        cypher_query_ = str(self._sanitize_parameters(cypher_query_, params))
+        try:
+            result = self.client.query(cypher_query_, params=params)
+        except Exception as e:
+            logger.error('Query: %s', cypher_query_)
+            logger.error('Parameters: %s', params)
+            logger.error('Error executing query: %s', e)
+            raise e
+
+        return result, None, None
+
+    def session(self, database: str | None = None) -> GraphDriverSession:
+        return NeptuneDriverSession(driver=self)
+
+    async def close(self) -> None:
+        return self.client.client.close()
+
+    async def _delete_all_data(self) -> Any:
+        return await self.execute_query('MATCH (n) DETACH DELETE n')
+
+    def delete_all_indexes(self) -> Coroutine[Any, Any, Any]:
+        return self.delete_all_indexes_impl()
+
+    async def delete_all_indexes_impl(self) -> Coroutine[Any, Any, Any]:
+        # No matter what happens above, always return True
+        return self.delete_aoss_indices()
+
+    async def create_aoss_indices(self):
+        for index in aoss_indices:
+            index_name = index['index_name']
+            client = self.aoss_client
+            if not client.indices.exists(index=index_name):
+                client.indices.create(index=index_name, body=index['body'])
+        # Sleep for 1 minute to let the index creation complete
+        await asyncio.sleep(60)
+
+    async def delete_aoss_indices(self):
+        for index in aoss_indices:
+            index_name = index['index_name']
+            client = self.aoss_client
+            if client.indices.exists(index=index_name):
+                client.indices.delete(index=index_name)
+
+    async def build_indices_and_constraints(self, delete_existing: bool = False):
+        # Neptune uses OpenSearch (AOSS) for indexing
+        if delete_existing:
+            await self.delete_aoss_indices()
+        await self.create_aoss_indices()
+
+    def run_aoss_query(self, name: str, query_text: str, limit: int = 10) -> dict[str, Any]:
+        for index in aoss_indices:
+            if name.lower() == index['index_name']:
+                index['query']['query']['multi_match']['query'] = query_text
+                query = {'size': limit, 'query': index['query']}
+                resp = self.aoss_client.search(body=query['query'], index=index['index_name'])
+                return resp
+        return {}
+
+    def save_to_aoss(self, name: str, data: list[dict]) -> int:
+        for index in aoss_indices:
+            if name.lower() == index['index_name']:
+                to_index = []
+                for d in data:
+                    item = {'_index': name, '_id': d['uuid']}
+                    for p in index['body']['mappings']['properties']:
+                        if p in d:
+                            item[p] = d[p]
+                    to_index.append(item)
+                success, failed = helpers.bulk(self.aoss_client, to_index, stats_only=True)
+                return success
+
+        return 0
+
+
+class NeptuneDriverSession(GraphDriverSession):
+    provider = GraphProvider.NEPTUNE
+
+    def __init__(self, driver: NeptuneDriver):  # type: ignore[reportUnknownArgumentType]
+        self.driver = driver
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Neptune, but method must exist
+        pass
+
+    async def close(self):
+        # No explicit close needed for Neptune, but method must exist
+        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):
+            res = None
+            for q in query:
+                res = await self.driver.execute_query(q, **kwargs)
+            return res
+        else:
+            return await self.driver.execute_query(str(query), **kwargs)

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