graphiti-core 0.21.0rc1__tar.gz → 0.21.0rc3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.21.0rc1
+Version: 0.21.0rc3
 Summary: A temporal graph building library
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
 Project-URL: Repository, https://github.com/getzep/graphiti

{graphiti_core-0.21.0rc1 → graphiti_core-0.21.0rc3}/graphiti_core/driver/driver.py

@@ -17,16 +17,19 @@ 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 OpenSearch, helpers
+    from opensearchpy import AsyncOpenSearch, helpers
 
     _HAS_OPENSEARCH = True
 except ImportError:
@@ -38,6 +41,13 @@ 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'
@@ -48,20 +58,19 @@ class GraphProvider(Enum):
 
 aoss_indices = [
     {
-        'index_name': 'entities',
+        'index_name': ENTITY_INDEX_NAME,
         'body': {
+            'settings': {'index': {'knn': True}},
             '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"},
+                    'group_id': {'type': 'keyword'},
+                    'created_at': {'type': 'date', 'format': 'strict_date_optional_time_nanos'},
                     'name_embedding': {
                         'type': 'knn_vector',
-                        'dims': EMBEDDING_DIM,
-                        'index': True,
-                        'similarity': 'cosine',
+                        'dimension': EMBEDDING_DIM,
                         'method': {
                             'engine': 'faiss',
                             'space_type': 'cosinesimil',
@@ -70,23 +79,23 @@ aoss_indices = [
                         },
                     },
                 }
-            }
+            },
         },
     },
     {
-        'index_name': 'communities',
+        'index_name': COMMUNITY_INDEX_NAME,
         'body': {
             'mappings': {
                 'properties': {
                     'uuid': {'type': 'keyword'},
                     'name': {'type': 'text'},
-                    'group_id': {'type': 'text'},
+                    'group_id': {'type': 'keyword'},
                 }
             }
         },
     },
     {
-        'index_name': 'episodes',
+        'index_name': EPISODE_INDEX_NAME,
         'body': {
             'mappings': {
                 'properties': {
@@ -94,31 +103,30 @@ aoss_indices = [
                     '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"},
+                    '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_edges',
+        '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': '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"},
+                    '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',
-                        'dims': EMBEDDING_DIM,
-                        'index': True,
-                        'similarity': 'cosine',
+                        'dimension': EMBEDDING_DIM,
                         'method': {
                             'engine': 'faiss',
                             'space_type': 'cosinesimil',
@@ -127,7 +135,7 @@ aoss_indices = [
                         },
                     },
                 }
-            }
+            },
         },
     },
 ]
@@ -163,7 +171,7 @@ class GraphDriver(ABC):
         ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
     )
     _database: str
-    aoss_client: OpenSearch | None  # type: ignore
+    aoss_client: AsyncOpenSearch | None  # type: ignore
 
     @abstractmethod
     def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
@@ -205,7 +213,7 @@ class GraphDriver(ABC):
             alias_name = index['index_name']
 
             # If alias already exists, skip (idempotent behavior)
-            if client.indices.exists_alias(name=alias_name):
+            if await client.indices.exists_alias(name=alias_name):
                 continue
 
             # Build a physical index name with timestamp
@@ -213,27 +221,67 @@ class GraphDriver(ABC):
             physical_index_name = f'{alias_name}_{ts_suffix}'
 
             # Create the index
-            client.indices.create(index=physical_index_name, body=index['body'])
+            await 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)
+            await client.indices.put_alias(index=physical_index_name, name=alias_name)
 
         # Allow some time for index creation
-        await asyncio.sleep(60)
+        await asyncio.sleep(1)
 
     async def delete_aoss_indices(self):
-        for index in aoss_indices:
-            index_name = index['index_name']
-            client = self.aoss_client
+        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
 
-            if not client:
-                logger.warning('No OpenSearch client found')
-                return
+                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}'")
 
-            if client.indices.exists(index=index_name):
-                client.indices.delete(index=index_name)
+            except Exception as e:
+                logger.error(f"Error deleting indices for alias '{alias_name}': {e}")
 
-    def save_to_aoss(self, name: str, data: list[dict]) -> int:
+    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')
@@ -243,16 +291,22 @@ class GraphDriver(ABC):
             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,
-                        '_routing': d.get('group_id'),  # shard routing
+                        '_id': d['uuid'],
+                        '_routing': d.get('group_id'),
+                        '_source': doc,
                     }
-                    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)
+                success, failed = await helpers.async_bulk(
+                    client, to_index, stats_only=True, request_timeout=60
+                )
 
                 return success if failed == 0 else success
 

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

@@ -28,7 +28,13 @@ logger = logging.getLogger(__name__)
 
 try:
     import boto3
-    from opensearchpy import OpenSearch, Urllib3AWSV4SignerAuth, Urllib3HttpConnection
+    from opensearchpy import (
+        AIOHttpConnection,
+        AsyncOpenSearch,
+        AWSV4SignerAuth,
+        Urllib3AWSV4SignerAuth,
+        Urllib3HttpConnection,
+    )
 
     _HAS_OPENSEARCH = True
 except ImportError:
@@ -50,6 +56,9 @@ class Neo4jDriver(GraphDriver):
         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(
@@ -61,15 +70,17 @@ class Neo4jDriver(GraphDriver):
         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
+                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}],
-                    http_auth=Urllib3AWSV4SignerAuth(  # type: ignore
-                        session.get_credentials(), session.region_name, 'aoss'
-                    ),
+                    auth=auth,
                     use_ssl=True,
                     verify_certs=True,
-                    connection_class=Urllib3HttpConnection,
+                    connection_class=AIOHttpConnection,
                     pool_maxsize=20,
                 )  # type: ignore
             except Exception as e:

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

@@ -237,12 +237,12 @@ class NeptuneDriver(GraphDriver):
                     '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):
-                client.indices.put_alias(index=index_name, name=alias_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)

{graphiti_core-0.21.0rc1 → graphiti_core-0.21.0rc3}/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):
@@ -256,13 +269,13 @@ class EntityEdge(Edge):
                 RETURN [x IN split(e.fact_embedding, ",") | toFloat(x)] as fact_embedding
             """
         elif driver.aoss_client:
-            resp = driver.aoss_client.search(
+            resp = await driver.aoss_client.search(
                 body={
                     'query': {'multi_match': {'query': self.uuid, 'fields': ['uuid']}},
                     'size': 1,
                 },
-                index='entity_edges',
-                routing=self.group_id,
+                index=ENTITY_EDGE_INDEX_NAME,
+                params={'routing': self.group_id},
             )
 
             if resp['hits']['hits']:
@@ -314,7 +327,7 @@ class EntityEdge(Edge):
             edge_data.update(self.attributes or {})
 
             if driver.aoss_client:
-                driver.save_to_aoss('entity_edges', [edge_data])  # pyright: ignore reportAttributeAccessIssue
+                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),
@@ -351,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-0.21.0rc1 → graphiti_core-0.21.0rc3}/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,

{graphiti_core-0.21.0rc1 → graphiti_core-0.21.0rc3}/graphiti_core/nodes.py

@@ -26,7 +26,14 @@ 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 (
+    COMMUNITY_INDEX_NAME,
+    ENTITY_EDGE_INDEX_NAME,
+    ENTITY_INDEX_NAME,
+    EPISODE_INDEX_NAME,
+    GraphDriver,
+    GraphProvider,
+)
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
 from graphiti_core.helpers import parse_db_date
@@ -94,13 +101,39 @@ class Node(BaseModel, ABC):
     async def delete(self, driver: GraphDriver):
         match driver.provider:
             case GraphProvider.NEO4J:
-                await driver.execute_query(
+                records, _, _ = await driver.execute_query(
                     """
-                    MATCH (n:Entity|Episodic|Community {uuid: $uuid})
+                    MATCH (n {uuid: $uuid})
+                    WHERE n:Entity OR n:Episodic OR n:Community
+                    OPTIONAL MATCH (n)-[r]-()
+                    WITH collect(r.uuid) AS edge_uuids, n
                     DETACH DELETE n
+                    RETURN edge_uuids
                     """,
                     uuid=self.uuid,
                 )
+
+                edge_uuids: list[str] = records[0].get('edge_uuids', []) if records else []
+
+                if driver.aoss_client:
+                    # Delete the node from OpenSearch indices
+                    for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME):
+                        await driver.aoss_client.delete(
+                            index=index,
+                            id=self.uuid,
+                            params={'routing': self.group_id},
+                        )
+
+                    # Bulk delete the detached edges
+                    if edge_uuids:
+                        actions = []
+                        for eid in edge_uuids:
+                            actions.append(
+                                {'delete': {'_index': ENTITY_EDGE_INDEX_NAME, '_id': eid}}
+                            )
+
+                        await driver.aoss_client.bulk(body=actions)
+
             case GraphProvider.KUZU:
                 for label in ['Episodic', 'Community']:
                     await driver.execute_query(
@@ -162,6 +195,32 @@ class Node(BaseModel, ABC):
                         group_id=group_id,
                         batch_size=batch_size,
                     )
+
+                if driver.aoss_client:
+                    await driver.aoss_client.delete_by_query(
+                        index=EPISODE_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
+                    await driver.aoss_client.delete_by_query(
+                        index=ENTITY_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
+                    await driver.aoss_client.delete_by_query(
+                        index=COMMUNITY_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
+                    await driver.aoss_client.delete_by_query(
+                        index=ENTITY_EDGE_INDEX_NAME,
+                        body={'query': {'term': {'group_id': group_id}}},
+                        params={'routing': group_id},
+                    )
+
             case GraphProvider.KUZU:
                 for label in ['Episodic', 'Community']:
                     await driver.execute_query(
@@ -240,6 +299,23 @@ class Node(BaseModel, ABC):
                 )
             case _:  # Neo4J, Neptune
                 async with driver.session() as session:
+                    # Collect all edge UUIDs before deleting nodes
+                    result = await session.run(
+                        """
+                        MATCH (n:Entity|Episodic|Community)
+                        WHERE n.uuid IN $uuids
+                        MATCH (n)-[r]-()
+                        RETURN collect(r.uuid) AS edge_uuids
+                        """,
+                        uuids=uuids,
+                    )
+
+                    record = await result.single()
+                    edge_uuids: list[str] = (
+                        record['edge_uuids'] if record and record['edge_uuids'] else []
+                    )
+
+                    # Now delete the nodes in batches
                     await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community)
@@ -253,6 +329,20 @@ class Node(BaseModel, ABC):
                         batch_size=batch_size,
                     )
 
+                if driver.aoss_client:
+                    for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME):
+                        await driver.aoss_client.delete_by_query(
+                            index=index,
+                            body={'query': {'terms': {'uuid': uuids}}},
+                        )
+
+                    if edge_uuids:
+                        actions = [
+                            {'delete': {'_index': ENTITY_EDGE_INDEX_NAME, '_id': eid}}
+                            for eid in edge_uuids
+                        ]
+                        await driver.aoss_client.bulk(body=actions)
+
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
 
@@ -286,7 +376,7 @@ class EpisodicNode(Node):
         }
 
         if driver.aoss_client:
-            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+            await driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
                 'episodes',
                 [episode_args],
             )
@@ -426,13 +516,13 @@ class EntityNode(Node):
                 RETURN [x IN split(n.name_embedding, ",") | toFloat(x)] as name_embedding
             """
         elif driver.aoss_client:
-            resp = driver.aoss_client.search(
+            resp = await driver.aoss_client.search(
                 body={
                     'query': {'multi_match': {'query': self.uuid, 'fields': ['uuid']}},
                     'size': 1,
                 },
-                index='entities',
-                routing=self.group_id,
+                index=ENTITY_INDEX_NAME,
+                params={'routing': self.group_id},
             )
 
             if resp['hits']['hits']:
@@ -479,7 +569,7 @@ class EntityNode(Node):
             labels = ':'.join(self.labels + ['Entity'])
 
             if driver.aoss_client:
-                driver.save_to_aoss('entities', [entity_data])  # pyright: ignore reportAttributeAccessIssue
+                await driver.save_to_aoss(ENTITY_INDEX_NAME, [entity_data])  # pyright: ignore reportAttributeAccessIssue
 
             result = await driver.execute_query(
                 get_entity_node_save_query(driver.provider, labels, bool(driver.aoss_client)),
@@ -577,7 +667,7 @@ class CommunityNode(Node):
 
     async def save(self, driver: GraphDriver):
         if driver.provider == GraphProvider.NEPTUNE:
-            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+            await driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
                 'communities',
                 [{'name': self.name, 'uuid': self.uuid, 'group_id': self.group_id}],
             )