graphiti-core 0.20.4__py3-none-any.whl → 0.21.0rc2__py3-none-any.whl

graphiti_core/driver/driver.py

@@ -14,15 +14,40 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import asyncio
 import copy
 import logging
+import os
 from abc import ABC, abstractmethod
 from collections.abc import Coroutine
+from datetime import datetime
 from enum import Enum
 from typing import Any
 
+from dotenv import load_dotenv
+
+from graphiti_core.embedder.client import EMBEDDING_DIM
+
+try:
+    from opensearchpy import AsyncOpenSearch, helpers
+
+    _HAS_OPENSEARCH = True
+except ImportError:
+    OpenSearch = None
+    helpers = None
+    _HAS_OPENSEARCH = False
+
 logger = logging.getLogger(__name__)
 
+DEFAULT_SIZE = 10
+
+load_dotenv()
+
+ENTITY_INDEX_NAME = os.environ.get('ENTITY_INDEX_NAME', 'entities')
+EPISODE_INDEX_NAME = os.environ.get('EPISODE_INDEX_NAME', 'episodes')
+COMMUNITY_INDEX_NAME = os.environ.get('COMMUNITY_INDEX_NAME', 'communities')
+ENTITY_EDGE_INDEX_NAME = os.environ.get('ENTITY_EDGE_INDEX_NAME', 'entity_edges')
+
 
 class GraphProvider(Enum):
     NEO4J = 'neo4j'
@@ -31,6 +56,91 @@ class GraphProvider(Enum):
     NEPTUNE = 'neptune'
 
 
+aoss_indices = [
+    {
+        'index_name': ENTITY_INDEX_NAME,
+        'body': {
+            'settings': {'index': {'knn': True}},
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'summary': {'type': 'text'},
+                    'group_id': {'type': 'keyword'},
+                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                    'name_embedding': {
+                        'type': 'knn_vector',
+                        'dimension': EMBEDDING_DIM,
+                        'method': {
+                            'engine': 'faiss',
+                            'space_type': 'cosinesimil',
+                            'name': 'hnsw',
+                            'parameters': {'ef_construction': 128, 'm': 16},
+                        },
+                    },
+                }
+            },
+        },
+    },
+    {
+        'index_name': COMMUNITY_INDEX_NAME,
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'group_id': {'type': 'keyword'},
+                }
+            }
+        },
+    },
+    {
+        'index_name': EPISODE_INDEX_NAME,
+        'body': {
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'content': {'type': 'text'},
+                    'source': {'type': 'text'},
+                    'source_description': {'type': 'text'},
+                    'group_id': {'type': 'keyword'},
+                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                    'valid_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                }
+            }
+        },
+    },
+    {
+        'index_name': ENTITY_EDGE_INDEX_NAME,
+        'body': {
+            'settings': {'index': {'knn': True}},
+            'mappings': {
+                'properties': {
+                    'uuid': {'type': 'keyword'},
+                    'name': {'type': 'text'},
+                    'fact': {'type': 'text'},
+                    'group_id': {'type': 'keyword'},
+                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                    'valid_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                    'expired_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                    'invalid_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
+                    'fact_embedding': {
+                        'type': 'knn_vector',
+                        'dimension': EMBEDDING_DIM,
+                        'method': {
+                            'engine': 'faiss',
+                            'space_type': 'cosinesimil',
+                            'name': 'hnsw',
+                            'parameters': {'ef_construction': 128, 'm': 16},
+                        },
+                    },
+                }
+            },
+        },
+    },
+]
+
+
 class GraphDriverSession(ABC):
     provider: GraphProvider
 
@@ -61,6 +171,7 @@ class GraphDriver(ABC):
         ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
     )
     _database: str
+    aoss_client: AsyncOpenSearch | None  # type: ignore
 
     @abstractmethod
     def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
@@ -87,3 +198,116 @@ 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 await 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
+            await client.indices.create(index=physical_index_name, body=index['body'])
+
+            # Point alias to it
+            await client.indices.put_alias(index=physical_index_name, name=alias_name)
+
+        # Allow some time for index creation
+        await asyncio.sleep(1)
+
+    async def delete_aoss_indices(self):
+        client = self.aoss_client
+
+        if not client:
+            logger.warning('No OpenSearch client found')
+            return
+
+        for entry in aoss_indices:
+            alias_name = entry['index_name']
+
+            try:
+                # Resolve alias → indices
+                alias_info = await client.indices.get_alias(name=alias_name)
+                indices = list(alias_info.keys())
+
+                if not indices:
+                    logger.info(f"No indices found for alias '{alias_name}'")
+                    continue
+
+                for index in indices:
+                    if await client.indices.exists(index=index):
+                        await client.indices.delete(index=index)
+                        logger.info(f"Deleted index '{index}' (alias: {alias_name})")
+                    else:
+                        logger.warning(f"Index '{index}' not found for alias '{alias_name}'")
+
+            except Exception as e:
+                logger.error(f"Error deleting indices for alias '{alias_name}': {e}")
+
+    async def clear_aoss_indices(self):
+        client = self.aoss_client
+
+        if not client:
+            logger.warning('No OpenSearch client found')
+            return
+
+        for index in aoss_indices:
+            index_name = index['index_name']
+
+            if await client.indices.exists(index=index_name):
+                try:
+                    # Delete all documents but keep the index
+                    response = await client.delete_by_query(
+                        index=index_name,
+                        body={'query': {'match_all': {}}},
+                    )
+                    logger.info(f"Cleared index '{index_name}': {response}")
+                except Exception as e:
+                    logger.error(f"Error clearing index '{index_name}': {e}")
+            else:
+                logger.warning(f"Index '{index_name}' does not exist")
+
+    async 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:
+                    doc = {}
+                    for p in index['body']['mappings']['properties']:
+                        if p in d:  # protect against missing fields
+                            doc[p] = d[p]
+
+                    item = {
+                        '_index': name,
+                        '_id': d['uuid'],
+                        '_routing': d.get('group_id'),
+                        '_source': doc,
+                    }
+                    to_index.append(item)
+
+                success, failed = await helpers.async_bulk(
+                    client, to_index, stats_only=True, request_timeout=60
+                )
+
+                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,44 @@ 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 (
+        AIOHttpConnection,
+        AsyncOpenSearch,
+        AWSV4SignerAuth,
+        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,
+        aws_profile_name: str | None = None,
+        aws_region: str | None = None,
+        aws_service: str | None = None,
+    ):
         super().__init__()
         self.client = AsyncGraphDatabase.driver(
             uri=uri,
@@ -37,6 +67,26 @@ class Neo4jDriver(GraphDriver):
         )
         self._database = database
 
+        self.aoss_client = None
+        if aoss_host and aoss_port and boto3 is not None:
+            try:
+                region = aws_region
+                service = aws_service
+                credentials = boto3.Session(profile_name=aws_profile_name).get_credentials()
+                auth = AWSV4SignerAuth(credentials, region or '', service or '')
+
+                self.aoss_client = AsyncOpenSearch(
+                    hosts=[{'host': aoss_host, 'port': aoss_port}],
+                    auth=auth,
+                    use_ssl=True,
+                    verify_certs=True,
+                    connection_class=AIOHttpConnection,
+                    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 +110,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'])
+                await 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):
+                await 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

@@ -25,7 +25,7 @@ from uuid import uuid4
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver, GraphProvider
+from graphiti_core.driver.driver import ENTITY_EDGE_INDEX_NAME, GraphDriver, GraphProvider
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
 from graphiti_core.helpers import parse_db_date
@@ -77,6 +77,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
@@ -108,6 +115,12 @@ 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):
@@ -255,6 +268,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 = 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 = """
@@ -292,14 +320,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:
+                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),
@@ -336,6 +364,35 @@ class EntityEdge(Edge):
             raise EdgeNotFoundError(uuid)
         return edges[0]
 
+    @classmethod
+    async def get_between_nodes(
+        cls, driver: GraphDriver, source_node_uuid: str, target_node_uuid: str
+    ):
+        match_query = """
+            MATCH (n:Entity {uuid: $source_node_uuid})-[e:RELATES_TO]->(m:Entity {uuid: $target_node_uuid})
+        """
+        if driver.provider == GraphProvider.KUZU:
+            match_query = """
+                MATCH (n:Entity {uuid: $source_node_uuid})
+                      -[:RELATES_TO]->(e:RelatesToNode_)
+                      -[:RELATES_TO]->(m:Entity {uuid: $target_node_uuid})
+            """
+
+        records, _, _ = await driver.execute_query(
+            match_query
+            + """
+            RETURN
+            """
+            + get_entity_edge_return_query(driver.provider),
+            source_node_uuid=source_node_uuid,
+            target_node_uuid=target_node_uuid,
+            routing_='r',
+        )
+
+        edges = [get_entity_edge_from_record(record, driver.provider) for record in records]
+
+        return edges
+
     @classmethod
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         if len(uuids) == 0:

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

@@ -60,9 +60,7 @@ from graphiti_core.search.search_config_recipes import (
 from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.search.search_utils import (
     RELEVANT_SCHEMA_LIMIT,
-    get_edge_invalidation_candidates,
     get_mentioned_nodes,
-    get_relevant_edges,
 )
 from graphiti_core.telemetry import capture_event
 from graphiti_core.utils.bulk_utils import (
@@ -1037,10 +1035,28 @@ class Graphiti:
 
         updated_edge = resolve_edge_pointers([edge], uuid_map)[0]
 
-        related_edges = (await get_relevant_edges(self.driver, [updated_edge], SearchFilters()))[0]
+        valid_edges = await EntityEdge.get_between_nodes(
+            self.driver, edge.source_node_uuid, edge.target_node_uuid
+        )
+
+        related_edges = (
+            await search(
+                self.clients,
+                updated_edge.fact,
+                group_ids=[updated_edge.group_id],
+                config=EDGE_HYBRID_SEARCH_RRF,
+                search_filter=SearchFilters(edge_uuids=[edge.uuid for edge in valid_edges]),
+            )
+        ).edges
         existing_edges = (
-            await get_edge_invalidation_candidates(self.driver, [updated_edge], SearchFilters())
-        )[0]
+            await search(
+                self.clients,
+                updated_edge.fact,
+                group_ids=[updated_edge.group_id],
+                config=EDGE_HYBRID_SEARCH_RRF,
+                search_filter=SearchFilters(),
+            )
+        ).edges
 
         resolved_edge, invalidated_edges, _ = await resolve_extracted_edge(
             self.llm_client,