graphiti-core 0.20.4__tar.gz → 0.21.0rc1__tar.gz

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.20.4
+Version: 0.21.0rc1
 Summary: A temporal graph building library
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
 Project-URL: Repository, https://github.com/getzep/graphiti
@@ -47,6 +47,9 @@ Provides-Extra: groq
 Requires-Dist: groq>=0.2.0; extra == 'groq'
 Provides-Extra: kuzu
 Requires-Dist: kuzu>=0.11.2; extra == 'kuzu'
+Provides-Extra: neo4j-opensearch
+Requires-Dist: boto3>=1.39.16; extra == 'neo4j-opensearch'
+Requires-Dist: opensearch-py>=3.0.0; extra == 'neo4j-opensearch'
 Provides-Extra: neptune
 Requires-Dist: boto3>=1.39.16; extra == 'neptune'
 Requires-Dist: langchain-aws>=0.2.29; extra == 'neptune'

graphiti_core-0.21.0rc1/graphiti_core/driver/driver.py

@@ -0,0 +1,259 @@
+"""
+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 copy
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from graphiti_core.embedder.client import EMBEDDING_DIM
+
+try:
+    from opensearchpy import OpenSearch, helpers
+
+    _HAS_OPENSEARCH = True
+except ImportError:
+    OpenSearch = None
+    helpers = None
+    _HAS_OPENSEARCH = False
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SIZE = 10
+
+
+class GraphProvider(Enum):
+    NEO4J = 'neo4j'
+    FALKORDB = 'falkordb'
+    KUZU = 'kuzu'
+    NEPTUNE = 'neptune'
+
+
+aoss_indices = [
+    {
+        'index_name': 'entities',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'summary': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                    'created_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                    'name_embedding': {
+                        'type': 'knn_vector',
+                        'dims': EMBEDDING_DIM,
+                        'index': True,
+                        'similarity': 'cosine',
+                        'method': {
+                            'engine': 'faiss',
+                            'space_type': 'cosinesimil',
+                            'name': 'hnsw',
+                            'parameters': {'ef_construction': 128, 'm': 16},
+                        },
+                    },
+                }
+            }
+        },
+    },
+    {
+        'index_name': 'communities',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                }
+            }
+        },
+    },
+    {
+        'index_name': 'episodes',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'content': {'type': 'text'},
+                    'source': {'type': 'text'},
+                    'source_description': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                    'created_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                    'valid_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                }
+            }
+        },
+    },
+    {
+        'index_name': 'entity_edges',
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'fact': {'type': 'text'},
+                    'group_id': {'type': 'text'},
+                    'created_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                    'valid_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                    'expired_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                    'invalid_at': {'type': 'date', 'format': "yyyy-MM-dd'T'HH:mm:ss.SSSZ"},
+                    'fact_embedding': {
+                        'type': 'knn_vector',
+                        'dims': EMBEDDING_DIM,
+                        'index': True,
+                        'similarity': 'cosine',
+                        'method': {
+                            'engine': 'faiss',
+                            'space_type': 'cosinesimil',
+                            'name': 'hnsw',
+                            'parameters': {'ef_construction': 128, 'm': 16},
+                        },
+                    },
+                }
+            }
+        },
+    },
+]
+
+
+class GraphDriverSession(ABC):
+    provider: GraphProvider
+
+    async def __aenter__(self):
+        return self
+
+    @abstractmethod
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Falkor, but method must exist
+        pass
+
+    @abstractmethod
+    async def run(self, query: str, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def execute_write(self, func, *args, **kwargs):
+        raise NotImplementedError()
+
+
+class GraphDriver(ABC):
+    provider: GraphProvider
+    fulltext_syntax: str = (
+        ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
+    )
+    _database: str
+    aoss_client: OpenSearch | None  # type: ignore
+
+    @abstractmethod
+    def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def session(self, database: str | None = None) -> GraphDriverSession:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def delete_all_indexes(self) -> Coroutine:
+        raise NotImplementedError()
+
+    def with_database(self, database: str) -> 'GraphDriver':
+        """
+        Returns a shallow copy of this driver with a different default database.
+        Reuses the same connection (e.g. FalkorDB, Neo4j).
+        """
+        cloned = copy.copy(self)
+        cloned._database = database
+
+        return cloned
+
+    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):
+        client = self.aoss_client
+        if not client:
+            logger.warning('No OpenSearch client found')
+            return
+
+        for index in aoss_indices:
+            alias_name = index['index_name']
+
+            # If alias already exists, skip (idempotent behavior)
+            if client.indices.exists_alias(name=alias_name):
+                continue
+
+            # Build a physical index name with timestamp
+            ts_suffix = datetime.utcnow().strftime('%Y%m%d%H%M%S')
+            physical_index_name = f'{alias_name}_{ts_suffix}'
+
+            # Create the index
+            client.indices.create(index=physical_index_name, body=index['body'])
+
+            # Point alias to it
+            client.indices.put_alias(index=physical_index_name, name=alias_name)
+
+        # Allow some time for index creation
+        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 not client:
+                logger.warning('No OpenSearch client found')
+                return
+
+            if client.indices.exists(index=index_name):
+                client.indices.delete(index=index_name)
+
+    def save_to_aoss(self, name: str, data: list[dict]) -> int:
+        client = self.aoss_client
+        if not client or not helpers:
+            logger.warning('No OpenSearch client found')
+            return 0
+
+        for index in aoss_indices:
+            if name.lower() == index['index_name']:
+                to_index = []
+                for d in data:
+                    item = {
+                        '_index': name,
+                        '_routing': d.get('group_id'),  # shard routing
+                    }
+                    for p in index['body']['mappings']['properties']:
+                        if p in d:  # protect against missing fields
+                            item[p] = d[p]
+                    to_index.append(item)
+
+                success, failed = helpers.bulk(client, to_index, stats_only=True)
+
+                return success if failed == 0 else success
+
+        return 0

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/driver/falkordb_driver.py

@@ -74,6 +74,7 @@ class FalkorDriverSession(GraphDriverSession):
 
 class FalkorDriver(GraphDriver):
     provider = GraphProvider.FALKORDB
+    aoss_client: None = None
 
     def __init__(
         self,

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/driver/kuzu_driver.py

@@ -92,6 +92,7 @@ SCHEMA_QUERIES = """
 
 class KuzuDriver(GraphDriver):
     provider: GraphProvider = GraphProvider.KUZU
+    aoss_client: None = None
 
     def __init__(
         self,

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/driver/neo4j_driver.py

@@ -22,14 +22,35 @@ from neo4j import AsyncGraphDatabase, EagerResult
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+from graphiti_core.helpers import semaphore_gather
 
 logger = logging.getLogger(__name__)
 
+try:
+    import boto3
+    from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection
+
+    _HAS_OPENSEARCH = True
+except ImportError:
+    boto3 = None
+    OpenSearch = None
+    Urllib3AWSV4SignerAuth = None
+    Urllib3HttpConnection = None
+    _HAS_OPENSEARCH = False
+
 
 class Neo4jDriver(GraphDriver):
     provider = GraphProvider.NEO4J
 
-    def __init__(self, uri: str, user: str | None, password: str | None, database: str = 'neo4j'):
+    def __init__(
+        self,
+        uri: str,
+        user: str | None,
+        password: str | None,
+        database: str = 'neo4j',
+        aoss_host: str | None = None,
+        aoss_port: int | None = None,
+    ):
         super().__init__()
         self.client = AsyncGraphDatabase.driver(
             uri=uri,
@@ -37,6 +58,24 @@ class Neo4jDriver(GraphDriver):
         )
         self._database = database
 
+        self.aoss_client = None
+        if aoss_host and aoss_port and boto3 is not None:
+            try:
+                session = boto3.Session()
+                self.aoss_client = OpenSearch(  # type: ignore
+                    hosts=[{'host': aoss_host, 'port': aoss_port}],
+                    http_auth=Urllib3AWSV4SignerAuth(  # type: ignore
+                        session.get_credentials(), session.region_name, 'aoss'
+                    ),
+                    use_ssl=True,
+                    verify_certs=True,
+                    connection_class=Urllib3HttpConnection,
+                    pool_maxsize=20,
+                )  # type: ignore
+            except Exception as e:
+                logger.warning(f'Failed to initialize OpenSearch client: {e}')
+                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
@@ -60,7 +99,14 @@ class Neo4jDriver(GraphDriver):
     async def close(self) -> None:
         return await self.client.close()
 
-    def delete_all_indexes(self) -> Coroutine[Any, Any, EagerResult]:
+    def delete_all_indexes(self) -> Coroutine:
+        if self.aoss_client:
+            return semaphore_gather(
+                self.client.execute_query(
+                    'CALL db.indexes() YIELD name DROP INDEX name',
+                ),
+                self.delete_aoss_indices(),
+            )
         return self.client.execute_query(
             'CALL db.indexes() YIELD name DROP INDEX name',
         )

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/driver/neptune_driver.py

@@ -22,16 +22,21 @@ from typing import Any
 
 import boto3
 from langchain_aws.graphs import NeptuneAnalyticsGraph, NeptuneGraph
-from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection, helpers
+from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection
 
-from graphiti_core.driver.driver import GraphDriver, GraphDriverSession, GraphProvider
+from graphiti_core.driver.driver import (
+    DEFAULT_SIZE,
+    GraphDriver,
+    GraphDriverSession,
+    GraphProvider,
+)
 
 logger = logging.getLogger(__name__)
-DEFAULT_SIZE = 10
 
-aoss_indices = [
+neptune_aoss_indices = [
     {
         'index_name': 'node_name_and_summary',
+        'alias_name': 'entities',
         'body': {
             'mappings': {
                 'properties': {
@@ -49,6 +54,7 @@ aoss_indices = [
     },
     {
         'index_name': 'community_name',
+        'alias_name': 'communities',
         'body': {
             'mappings': {
                 'properties': {
@@ -65,6 +71,7 @@ aoss_indices = [
     },
     {
         'index_name': 'episode_content',
+        'alias_name': 'episodes',
         'body': {
             'mappings': {
                 'properties': {
@@ -88,6 +95,7 @@ aoss_indices = [
     },
     {
         'index_name': 'edge_name_and_fact',
+        'alias_name': 'facts',
         'body': {
             'mappings': {
                 'properties': {
@@ -220,54 +228,27 @@ class NeptuneDriver(GraphDriver):
     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:
+        for index in neptune_aoss_indices:
             index_name = index['index_name']
             client = self.aoss_client
+            if not client:
+                raise ValueError(
+                    'You must provide an AOSS endpoint to create an OpenSearch driver.'
+                )
             if not client.indices.exists(index=index_name):
                 client.indices.create(index=index_name, body=index['body'])
+
+            alias_name = index.get('alias_name', index_name)
+
+            if not client.indices.exists_alias(name=alias_name, index=index_name):
+                client.indices.put_alias(index=index_name, name=alias_name)
+
         # 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)
-
-    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}
-                    for p in index['body']['mappings']['properties']:
-                        item[p] = d[p]
-                    to_index.append(item)
-                success, failed = helpers.bulk(self.aoss_client, to_index, stats_only=True)
-                if failed > 0:
-                    return success
-                else:
-                    return 0
-
-        return 0
+    def delete_all_indexes(self) -> Coroutine[Any, Any, Any]:
+        return self.delete_all_indexes_impl()
 
 
 class NeptuneDriverSession(GraphDriverSession):

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/edges.py

@@ -255,6 +255,21 @@ 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 = driver.aoss_client.search(
+                body={
+                    'query': {'multi_match': {'query': self.uuid, 'fields': ['uuid']}},
+                    'size': 1,
+                },
+                index='entity_edges',
+                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 = """
@@ -292,14 +307,14 @@ 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),
+                get_entity_edge_save_query(driver.provider, has_aoss=bool(driver.aoss_client)),
                 **edge_data,
             )
         else:
             edge_data.update(self.attributes or {})
 
-            if driver.provider == GraphProvider.NEPTUNE:
-                driver.save_to_aoss('edge_name_and_fact', [edge_data])  # pyright: ignore reportAttributeAccessIssue
+            if driver.aoss_client:
+                driver.save_to_aoss('entity_edges', [edge_data])  # pyright: ignore reportAttributeAccessIssue
 
             result = await driver.execute_query(
                 get_entity_edge_save_query(driver.provider),

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/embedder/client.py

@@ -14,12 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 
 from pydantic import BaseModel, Field
 
-EMBEDDING_DIM = 1024
+EMBEDDING_DIM = int(os.getenv('EMBEDDING_DIM', 1024))
 
 
 class EmbedderConfig(BaseModel):

{graphiti_core-0.20.4 → graphiti_core-0.21.0rc1}/graphiti_core/models/edges/edge_db_queries.py

@@ -60,7 +60,7 @@ EPISODIC_EDGE_RETURN = """
 """
 
 
-def get_entity_edge_save_query(provider: GraphProvider) -> str:
+def get_entity_edge_save_query(provider: GraphProvider, has_aoss: bool = False) -> str:
     match provider:
         case GraphProvider.FALKORDB:
             return """
@@ -99,17 +99,28 @@ def get_entity_edge_save_query(provider: GraphProvider) -> str:
                 RETURN e.uuid AS uuid
             """
         case _:  # Neo4j
-            return """
-                MATCH (source:Entity {uuid: $edge_data.source_uuid})
-                MATCH (target:Entity {uuid: $edge_data.target_uuid})
-                MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
-                SET e = $edge_data
-                WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)
+            save_embedding_query = (
+                """WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)"""
+                if not has_aoss
+                else ''
+            )
+            return (
+                (
+                    """
+                        MATCH (source:Entity {uuid: $edge_data.source_uuid})
+                        MATCH (target:Entity {uuid: $edge_data.target_uuid})
+                        MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+                        SET e = $edge_data
+                        """
+                    + save_embedding_query
+                )
+                + """
                 RETURN e.uuid AS uuid
-            """
+                """
+            )
 
 
-def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
+def get_entity_edge_save_bulk_query(provider: GraphProvider, has_aoss: bool = False) -> str:
     match provider:
         case GraphProvider.FALKORDB:
             return """
@@ -152,15 +163,24 @@ def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
                 RETURN e.uuid AS uuid
             """
         case _:
-            return """
-                UNWIND $entity_edges AS edge
-                MATCH (source:Entity {uuid: edge.source_node_uuid})
-                MATCH (target:Entity {uuid: edge.target_node_uuid})
-                MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
-                SET e = edge
-                WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)
+            save_embedding_query = (
+                'WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                """
+                    UNWIND $entity_edges AS edge
+                    MATCH (source:Entity {uuid: edge.source_node_uuid})
+                    MATCH (target:Entity {uuid: edge.target_node_uuid})
+                    MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
+                    SET e = edge
+                    """
+                + save_embedding_query
+                + """
                 RETURN edge.uuid AS uuid
             """
+            )
 
 
 def get_entity_edge_return_query(provider: GraphProvider) -> str: