graphiti-core 0.20.3__py3-none-any.whl → 0.21.0rc1__py3-none-any.whl

graphiti_core/driver/driver.py

@@ -14,15 +14,30 @@ 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'
@@ -31,6 +46,93 @@ class GraphProvider(Enum):
     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
 
@@ -61,6 +163,7 @@ class GraphDriver(ABC):
         ''  # 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:
@@ -87,3 +190,70 @@ class GraphDriver(ABC):
         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/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/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/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/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/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/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/graphiti.py

@@ -122,6 +122,11 @@ class AddBulkEpisodeResults(BaseModel):
     community_edges: list[CommunityEdge]
 
 
+class AddTripletResults(BaseModel):
+    nodes: list[EntityNode]
+    edges: list[EntityEdge]
+
+
 class Graphiti:
     def __init__(
         self,
@@ -1015,7 +1020,9 @@ class Graphiti:
 
         return SearchResults(edges=edges, nodes=nodes)
 
-    async def add_triplet(self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode):
+    async def add_triplet(
+        self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode
+    ) -> AddTripletResults:
         if source_node.name_embedding is None:
             await source_node.generate_name_embedding(self.embedder)
         if target_node.name_embedding is None:
@@ -1059,6 +1066,7 @@ class Graphiti:
         await create_entity_node_embeddings(self.embedder, nodes)
 
         await add_nodes_and_edges_bulk(self.driver, [], [], nodes, edges, self.embedder)
+        return AddTripletResults(edges=edges, nodes=nodes)
 
     async def remove_episode(self, episode_uuid: str):
         # Find the episode to be deleted

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:

graphiti_core/models/nodes/node_db_queries.py

@@ -126,7 +126,7 @@ EPISODIC_NODE_RETURN_NEPTUNE = """
 """
 
 
-def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
+def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: bool = False) -> str:
     match provider:
         case GraphProvider.FALKORDB:
             return f"""
@@ -161,16 +161,27 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
                 RETURN n.uuid AS uuid
             """
         case _:
-            return f"""
+            save_embedding_query = (
+                'WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                f"""
                 MERGE (n:Entity {{uuid: $entity_data.uuid}})
                 SET n:{labels}
                 SET n = $entity_data
-                WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)
+                """
+                + save_embedding_query
+                + """
                 RETURN n.uuid AS uuid
             """
+            )
 
 
-def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict]) -> str | Any:
+def get_entity_node_save_bulk_query(
+    provider: GraphProvider, nodes: list[dict], has_aoss: bool = False
+) -> str | Any:
     match provider:
         case GraphProvider.FALKORDB:
             queries = []
@@ -222,14 +233,23 @@ def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict])
                 RETURN n.uuid AS uuid
             """
         case _:  # Neo4j
-            return """
-                UNWIND $nodes AS node
-                MERGE (n:Entity {uuid: node.uuid})
-                SET n:$(node.labels)
-                SET n = node
-                WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
+            save_embedding_query = (
+                'WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                """
+                    UNWIND $nodes AS node
+                    MERGE (n:Entity {uuid: node.uuid})
+                    SET n:$(node.labels)
+                    SET n = node
+                    """
+                + save_embedding_query
+                + """
                 RETURN n.uuid AS uuid
             """
+            )
 
 
 def get_entity_node_return_query(provider: GraphProvider) -> str: