graphiti-core 0.18.9__py3-none-any.whl → 0.19.0rc2__py3-none-any.whl

graphiti_core/nodes.py

@@ -31,11 +31,13 @@ from graphiti_core.errors import NodeNotFoundError
 from graphiti_core.helpers import parse_db_date
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_RETURN,
+    COMMUNITY_NODE_RETURN_NEPTUNE,
     ENTITY_NODE_RETURN,
     EPISODIC_NODE_RETURN,
-    EPISODIC_NODE_SAVE,
+    EPISODIC_NODE_RETURN_NEPTUNE,
     get_community_node_save_query,
     get_entity_node_save_query,
+    get_episode_node_save_query,
 )
 from graphiti_core.utils.datetime_utils import utc_now
 
@@ -89,23 +91,24 @@ class Node(BaseModel, ABC):
     async def save(self, driver: GraphDriver): ...
 
     async def delete(self, driver: GraphDriver):
-        if driver.provider == GraphProvider.FALKORDB:
-            for label in ['Entity', 'Episodic', 'Community']:
+        match driver.provider:
+            case GraphProvider.NEO4J:
                 await driver.execute_query(
-                    f"""
-                    MATCH (n:{label} {{uuid: $uuid}})
-                    DETACH DELETE n
-                    """,
-                    uuid=self.uuid,
-                )
-        else:
-            await driver.execute_query(
-                """
+                    """
                 MATCH (n:Entity|Episodic|Community {uuid: $uuid})
                 DETACH DELETE n
                 """,
-                uuid=self.uuid,
-            )
+                    uuid=self.uuid,
+                )
+            case _:  # FalkorDB and Neptune
+                for label in ['Entity', 'Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label} {{uuid: $uuid}})
+                        DETACH DELETE n
+                        """,
+                        uuid=self.uuid,
+                    )
 
         logger.debug(f'Deleted Node: {self.uuid}')
 
@@ -119,28 +122,30 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_group_id(cls, driver: GraphDriver, group_id: str, batch_size: int = 100):
-        if driver.provider == GraphProvider.FALKORDB:
-            for label in ['Entity', 'Episodic', 'Community']:
-                await driver.execute_query(
-                    f"""
-                    MATCH (n:{label} {{group_id: $group_id}})
-                    DETACH DELETE n
-                    """,
-                    group_id=group_id,
-                )
-        else:
-            async with driver.session() as session:
-                await session.run(
-                    """
-                    MATCH (n:Entity|Episodic|Community {group_id: $group_id})
-                    CALL {
-                        WITH n
+        match driver.provider:
+            case GraphProvider.NEO4J:
+                async with driver.session() as session:
+                    await session.run(
+                        """
+                        MATCH (n:Entity|Episodic|Community {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            DETACH DELETE n
+                        } IN TRANSACTIONS OF $batch_size ROWS
+                        """,
+                        group_id=group_id,
+                        batch_size=batch_size,
+                    )
+
+            case _:  # FalkorDB and Neptune
+                for label in ['Entity', 'Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label} {{group_id: $group_id}})
                         DETACH DELETE n
-                    } IN TRANSACTIONS OF $batch_size ROWS
-                    """,
-                    group_id=group_id,
-                    batch_size=batch_size,
-                )
+                        """,
+                        group_id=group_id,
+                    )
 
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str], batch_size: int = 100):
@@ -189,11 +194,25 @@ class EpisodicNode(Node):
     )
 
     async def save(self, driver: GraphDriver):
+        if driver.provider == GraphProvider.NEPTUNE:
+            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+                'episode_content',
+                [
+                    {
+                        'uuid': self.uuid,
+                        'group_id': self.group_id,
+                        'source': self.source.value,
+                        'content': self.content,
+                        'source_description': self.source_description,
+                    }
+                ],
+            )
         result = await driver.execute_query(
-            EPISODIC_NODE_SAVE,
+            get_episode_node_save_query(driver.provider),
             uuid=self.uuid,
             name=self.name,
             group_id=self.group_id,
+            group_label='Episodic_' + self.group_id.replace('-', ''),
             source_description=self.source_description,
             content=self.content,
             entity_edges=self.entity_edges,
@@ -213,7 +232,11 @@ class EpisodicNode(Node):
             MATCH (e:Episodic {uuid: $uuid})
             RETURN
             """
-            + EPISODIC_NODE_RETURN,
+            + (
+                EPISODIC_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else EPISODIC_NODE_RETURN
+            ),
             uuid=uuid,
             routing_='r',
         )
@@ -233,7 +256,11 @@ class EpisodicNode(Node):
             WHERE e.uuid IN $uuids
             RETURN DISTINCT
             """
-            + EPISODIC_NODE_RETURN,
+            + (
+                EPISODIC_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else EPISODIC_NODE_RETURN
+            ),
             uuids=uuids,
             routing_='r',
         )
@@ -262,7 +289,11 @@ class EpisodicNode(Node):
             + """
             RETURN DISTINCT
             """
-            + EPISODIC_NODE_RETURN
+            + (
+                EPISODIC_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else EPISODIC_NODE_RETURN
+            )
             + """
             ORDER BY uuid DESC
             """
@@ -284,7 +315,11 @@ class EpisodicNode(Node):
             MATCH (e:Episodic)-[r:MENTIONS]->(n:Entity {uuid: $entity_node_uuid})
             RETURN DISTINCT
             """
-            + EPISODIC_NODE_RETURN,
+            + (
+                EPISODIC_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else EPISODIC_NODE_RETURN
+            ),
             entity_node_uuid=entity_node_uuid,
             routing_='r',
         )
@@ -311,11 +346,18 @@ class EntityNode(Node):
         return self.name_embedding
 
     async def load_name_embedding(self, driver: GraphDriver):
-        records, _, _ = await driver.execute_query(
+        if driver.provider == GraphProvider.NEPTUNE:
+            query: LiteralString = """
+                MATCH (n:Entity {uuid: $uuid})
+                RETURN [x IN split(n.name_embedding, ",") | toFloat(x)] as name_embedding
             """
-            MATCH (n:Entity {uuid: $uuid})
-            RETURN n.name_embedding AS name_embedding
-            """,
+        else:
+            query: LiteralString = """
+                MATCH (n:Entity {uuid: $uuid})
+                RETURN n.name_embedding AS name_embedding
+            """
+        records, _, _ = await driver.execute_query(
+            query,
             uuid=self.uuid,
             routing_='r',
         )
@@ -336,7 +378,10 @@ class EntityNode(Node):
         }
         entity_data.update(self.attributes or {})
 
-        labels = ':'.join(self.labels + ['Entity'])
+        if driver.provider == GraphProvider.NEPTUNE:
+            driver.save_to_aoss('node_name_and_summary', [entity_data])  # pyright: ignore reportAttributeAccessIssue
+
+        labels = ':'.join(self.labels + ['Entity', 'Entity_' + self.group_id.replace('-', '')])
 
         result = await driver.execute_query(
             get_entity_node_save_query(driver.provider, labels),
@@ -433,8 +478,13 @@ class CommunityNode(Node):
     summary: str = Field(description='region summary of member nodes', default_factory=str)
 
     async def save(self, driver: GraphDriver):
+        if driver.provider == GraphProvider.NEPTUNE:
+            driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
+                'community_name',
+                [{'name': self.name, 'uuid': self.uuid, 'group_id': self.group_id}],
+            )
         result = await driver.execute_query(
-            get_community_node_save_query(driver.provider),
+            get_community_node_save_query(driver.provider),  # type: ignore
             uuid=self.uuid,
             name=self.name,
             group_id=self.group_id,
@@ -457,11 +507,19 @@ class CommunityNode(Node):
         return self.name_embedding
 
     async def load_name_embedding(self, driver: GraphDriver):
-        records, _, _ = await driver.execute_query(
+        if driver.provider == GraphProvider.NEPTUNE:
+            query: LiteralString = """
+                MATCH (c:Community {uuid: $uuid})
+                RETURN [x IN split(c.name_embedding, ",") | toFloat(x)] as name_embedding
             """
+        else:
+            query: LiteralString = """
             MATCH (c:Community {uuid: $uuid})
             RETURN c.name_embedding AS name_embedding
-            """,
+            """
+
+        records, _, _ = await driver.execute_query(
+            query,
             uuid=self.uuid,
             routing_='r',
         )
@@ -478,7 +536,11 @@ class CommunityNode(Node):
             MATCH (n:Community {uuid: $uuid})
             RETURN
             """
-            + COMMUNITY_NODE_RETURN,
+            + (
+                COMMUNITY_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else COMMUNITY_NODE_RETURN
+            ),
             uuid=uuid,
             routing_='r',
         )
@@ -498,7 +560,11 @@ class CommunityNode(Node):
             WHERE n.uuid IN $uuids
             RETURN
             """
-            + COMMUNITY_NODE_RETURN,
+            + (
+                COMMUNITY_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else COMMUNITY_NODE_RETURN
+            ),
             uuids=uuids,
             routing_='r',
         )
@@ -527,7 +593,11 @@ class CommunityNode(Node):
             + """
             RETURN
             """
-            + COMMUNITY_NODE_RETURN
+            + (
+                COMMUNITY_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else COMMUNITY_NODE_RETURN
+            )
             + """
             ORDER BY n.uuid DESC
             """

graphiti_core/search/search.py

@@ -21,6 +21,7 @@ from time import time
 from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.edges import EntityEdge
+from graphiti_core.embedder.client import EMBEDDING_DIM
 from graphiti_core.errors import SearchRerankerError
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import semaphore_gather
@@ -29,6 +30,7 @@ from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
     CommunityReranker,
     CommunitySearchConfig,
+    CommunitySearchMethod,
     EdgeReranker,
     EdgeSearchConfig,
     EdgeSearchMethod,
@@ -81,11 +83,29 @@ async def search(
 
     if query.strip() == '':
         return SearchResults()
-    query_vector = (
-        query_vector
-        if query_vector is not None
-        else await embedder.create(input_data=[query.replace('\n', ' ')])
-    )
+
+    if (
+        config.edge_config
+        and EdgeSearchMethod.cosine_similarity in config.edge_config.search_methods
+        or config.edge_config
+        and EdgeReranker.mmr == config.edge_config.reranker
+        or config.node_config
+        and NodeSearchMethod.cosine_similarity in config.node_config.search_methods
+        or config.node_config
+        and NodeReranker.mmr == config.node_config.reranker
+        or (
+            config.community_config
+            and CommunitySearchMethod.cosine_similarity in config.community_config.search_methods
+        )
+        or (config.community_config and CommunityReranker.mmr == config.community_config.reranker)
+    ):
+        search_vector = (
+            query_vector
+            if query_vector is not None
+            else await embedder.create(input_data=[query.replace('\n', ' ')])
+        )
+    else:
+        search_vector = [0.0] * EMBEDDING_DIM
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids and group_ids != [''] else None
@@ -99,7 +119,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.edge_config,
             search_filter,
@@ -112,7 +132,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.node_config,
             search_filter,
@@ -125,7 +145,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.episode_config,
             search_filter,
@@ -136,7 +156,7 @@ async def search(
             driver,
             cross_encoder,
             query,
-            query_vector,
+            search_vector,
             group_ids,
             config.community_config,
             config.limit,