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

graphiti_core/nodes.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
@@ -31,11 +32,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,
-    ENTITY_NODE_RETURN,
+    COMMUNITY_NODE_RETURN_NEPTUNE,
     EPISODIC_NODE_RETURN,
-    EPISODIC_NODE_SAVE,
+    EPISODIC_NODE_RETURN_NEPTUNE,
     get_community_node_save_query,
+    get_entity_node_return_query,
     get_entity_node_save_query,
+    get_episode_node_save_query,
 )
 from graphiti_core.utils.datetime_utils import utc_now
 
@@ -89,23 +92,49 @@ 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}})
+                    """
+                    MATCH (n:Entity|Episodic|Community {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,
-            )
+            case GraphProvider.KUZU:
+                for label in ['Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label} {{uuid: $uuid}})
+                        DETACH DELETE n
+                        """,
+                        uuid=self.uuid,
+                    )
+                # Entity edges are actually nodes in Kuzu, so simple `DETACH DELETE` will not work.
+                # Explicitly delete the "edge" nodes first, then the entity node.
+                await driver.execute_query(
+                    """
+                    MATCH (n:Entity {uuid: $uuid})-[:RELATES_TO]->(e:RelatesToNode_)
+                    DETACH DELETE e
+                    """,
+                    uuid=self.uuid,
+                )
+                await driver.execute_query(
+                    """
+                    MATCH (n:Entity {uuid: $uuid})
+                    DETACH DELETE n
+                    """,
+                    uuid=self.uuid,
+                )
+            case _:  # FalkorDB, 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,55 +148,110 @@ 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']:
+        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 GraphProvider.KUZU:
+                for label in ['Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label} {{group_id: $group_id}})
+                        DETACH DELETE n
+                        """,
+                        group_id=group_id,
+                    )
+                # Entity edges are actually nodes in Kuzu, so simple `DETACH DELETE` will not work.
+                # Explicitly delete the "edge" nodes first, then the entity node.
                 await driver.execute_query(
-                    f"""
-                    MATCH (n:{label} {{group_id: $group_id}})
-                    DETACH DELETE n
+                    """
+                    MATCH (n:Entity {group_id: $group_id})-[:RELATES_TO]->(e:RelatesToNode_)
+                    DETACH DELETE e
                     """,
                     group_id=group_id,
                 )
-        else:
-            async with driver.session() as session:
-                await session.run(
+                await driver.execute_query(
                     """
-                    MATCH (n:Entity|Episodic|Community {group_id: $group_id})
-                    CALL {
-                        WITH n
-                        DETACH DELETE n
-                    } IN TRANSACTIONS OF $batch_size ROWS
+                    MATCH (n:Entity {group_id: $group_id})
+                    DETACH DELETE n
                     """,
                     group_id=group_id,
-                    batch_size=batch_size,
                 )
+            case _:  # FalkorDB, Neptune
+                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,
+                    )
 
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str], batch_size: int = 100):
-        if driver.provider == GraphProvider.FALKORDB:
-            for label in ['Entity', 'Episodic', 'Community']:
+        match driver.provider:
+            case GraphProvider.FALKORDB:
+                for label in ['Entity', 'Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label})
+                        WHERE n.uuid IN $uuids
+                        DETACH DELETE n
+                        """,
+                        uuids=uuids,
+                    )
+            case GraphProvider.KUZU:
+                for label in ['Episodic', 'Community']:
+                    await driver.execute_query(
+                        f"""
+                        MATCH (n:{label})
+                        WHERE n.uuid IN $uuids
+                        DETACH DELETE n
+                        """,
+                        uuids=uuids,
+                    )
+                # Entity edges are actually nodes in Kuzu, so simple `DETACH DELETE` will not work.
+                # Explicitly delete the "edge" nodes first, then the entity node.
                 await driver.execute_query(
-                    f"""
-                       MATCH (n:{label})
-                       WHERE n.uuid IN $uuids
-                       DETACH DELETE n
-                       """,
+                    """
+                    MATCH (n:Entity)-[:RELATES_TO]->(e:RelatesToNode_)
+                    WHERE n.uuid IN $uuids
+                    DETACH DELETE e
+                    """,
                     uuids=uuids,
                 )
-        else:
-            async with driver.session() as session:
-                await session.run(
+                await driver.execute_query(
                     """
-                    MATCH (n:Entity|Episodic|Community)
+                    MATCH (n:Entity)
                     WHERE n.uuid IN $uuids
-                    CALL {
-                        WITH n
-                        DETACH DELETE n
-                    } IN TRANSACTIONS OF $batch_size ROWS
+                    DETACH DELETE n
                     """,
                     uuids=uuids,
-                    batch_size=batch_size,
                 )
+            case _:  # Neo4J, Neptune
+                async with driver.session() as session:
+                    await session.run(
+                        """
+                        MATCH (n:Entity|Episodic|Community)
+                        WHERE n.uuid IN $uuids
+                        CALL {
+                            WITH n
+                            DETACH DELETE n
+                        } IN TRANSACTIONS OF $batch_size ROWS
+                        """,
+                        uuids=uuids,
+                        batch_size=batch_size,
+                    )
 
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
@@ -189,17 +273,37 @@ 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,
+                    }
+                ],
+            )
+
+        episode_args = {
+            'uuid': self.uuid,
+            'name': self.name,
+            'group_id': self.group_id,
+            'source_description': self.source_description,
+            'content': self.content,
+            'entity_edges': self.entity_edges,
+            'created_at': self.created_at,
+            'valid_at': self.valid_at,
+            'source': self.source.value,
+        }
+
+        if driver.provider == GraphProvider.NEO4J:
+            episode_args['group_label'] = 'Episodic_' + self.group_id.replace('-', '')
+
         result = await driver.execute_query(
-            EPISODIC_NODE_SAVE,
-            uuid=self.uuid,
-            name=self.name,
-            group_id=self.group_id,
-            source_description=self.source_description,
-            content=self.content,
-            entity_edges=self.entity_edges,
-            created_at=self.created_at,
-            valid_at=self.valid_at,
-            source=self.source.value,
+            get_episode_node_save_query(driver.provider), **episode_args
         )
 
         logger.debug(f'Saved Node to Graph: {self.uuid}')
@@ -213,7 +317,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 +341,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 +374,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 +400,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 +431,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',
         )
@@ -334,14 +461,25 @@ class EntityNode(Node):
             'summary': self.summary,
             'created_at': self.created_at,
         }
-        entity_data.update(self.attributes or {})
 
-        labels = ':'.join(self.labels + ['Entity'])
+        if driver.provider == GraphProvider.KUZU:
+            entity_data['attributes'] = json.dumps(self.attributes)
+            entity_data['labels'] = list(set(self.labels + ['Entity']))
+            result = await driver.execute_query(
+                get_entity_node_save_query(driver.provider, labels=''),
+                **entity_data,
+            )
+        else:
+            entity_data.update(self.attributes or {})
+            labels = ':'.join(self.labels + ['Entity', 'Entity_' + self.group_id.replace('-', '')])
 
-        result = await driver.execute_query(
-            get_entity_node_save_query(driver.provider, labels),
-            entity_data=entity_data,
-        )
+            if driver.provider == GraphProvider.NEPTUNE:
+                driver.save_to_aoss('node_name_and_summary', [entity_data])  # pyright: ignore reportAttributeAccessIssue
+
+            result = await driver.execute_query(
+                get_entity_node_save_query(driver.provider, labels),
+                entity_data=entity_data,
+            )
 
         logger.debug(f'Saved Node to Graph: {self.uuid}')
 
@@ -354,12 +492,12 @@ class EntityNode(Node):
             MATCH (n:Entity {uuid: $uuid})
             RETURN
             """
-            + ENTITY_NODE_RETURN,
+            + get_entity_node_return_query(driver.provider),
             uuid=uuid,
             routing_='r',
         )
 
-        nodes = [get_entity_node_from_record(record) for record in records]
+        nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
         if len(nodes) == 0:
             raise NodeNotFoundError(uuid)
@@ -374,12 +512,12 @@ class EntityNode(Node):
             WHERE n.uuid IN $uuids
             RETURN
             """
-            + ENTITY_NODE_RETURN,
+            + get_entity_node_return_query(driver.provider),
             uuids=uuids,
             routing_='r',
         )
 
-        nodes = [get_entity_node_from_record(record) for record in records]
+        nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
         return nodes
 
@@ -411,7 +549,7 @@ class EntityNode(Node):
             + """
             RETURN
             """
-            + ENTITY_NODE_RETURN
+            + get_entity_node_return_query(driver.provider)
             + with_embeddings_query
             + """
             ORDER BY n.uuid DESC
@@ -423,7 +561,7 @@ class EntityNode(Node):
             routing_='r',
         )
 
-        nodes = [get_entity_node_from_record(record) for record in records]
+        nodes = [get_entity_node_from_record(record, driver.provider) for record in records]
 
         return nodes
 
@@ -433,8 +571,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 +600,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',
         )
@@ -475,10 +626,14 @@ class CommunityNode(Node):
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
-            MATCH (n:Community {uuid: $uuid})
+            MATCH (c:Community {uuid: $uuid})
             RETURN
             """
-            + COMMUNITY_NODE_RETURN,
+            + (
+                COMMUNITY_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else COMMUNITY_NODE_RETURN
+            ),
             uuid=uuid,
             routing_='r',
         )
@@ -494,11 +649,15 @@ class CommunityNode(Node):
     async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
-            MATCH (n:Community)
-            WHERE n.uuid IN $uuids
+            MATCH (c:Community)
+            WHERE c.uuid IN $uuids
             RETURN
             """
-            + COMMUNITY_NODE_RETURN,
+            + (
+                COMMUNITY_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else COMMUNITY_NODE_RETURN
+            ),
             uuids=uuids,
             routing_='r',
         )
@@ -515,21 +674,25 @@ class CommunityNode(Node):
         limit: int | None = None,
         uuid_cursor: str | None = None,
     ):
-        cursor_query: LiteralString = 'AND n.uuid < $uuid' if uuid_cursor else ''
+        cursor_query: LiteralString = 'AND c.uuid < $uuid' if uuid_cursor else ''
         limit_query: LiteralString = 'LIMIT $limit' if limit is not None else ''
 
         records, _, _ = await driver.execute_query(
             """
-            MATCH (n:Community)
-            WHERE n.group_id IN $group_ids
+            MATCH (c:Community)
+            WHERE c.group_id IN $group_ids
             """
             + cursor_query
             + """
             RETURN
             """
-            + COMMUNITY_NODE_RETURN
+            + (
+                COMMUNITY_NODE_RETURN_NEPTUNE
+                if driver.provider == GraphProvider.NEPTUNE
+                else COMMUNITY_NODE_RETURN
+            )
             + """
-            ORDER BY n.uuid DESC
+            ORDER BY c.uuid DESC
             """
             + limit_query,
             group_ids=group_ids,
@@ -566,25 +729,35 @@ def get_episodic_node_from_record(record: Any) -> EpisodicNode:
     )
 
 
-def get_entity_node_from_record(record: Any) -> EntityNode:
+def get_entity_node_from_record(record: Any, provider: GraphProvider) -> EntityNode:
+    if provider == GraphProvider.KUZU:
+        attributes = json.loads(record['attributes']) if record['attributes'] else {}
+    else:
+        attributes = record['attributes']
+        attributes.pop('uuid', None)
+        attributes.pop('name', None)
+        attributes.pop('group_id', None)
+        attributes.pop('name_embedding', None)
+        attributes.pop('summary', None)
+        attributes.pop('created_at', None)
+        attributes.pop('labels', None)
+
+    labels = record.get('labels', [])
+    group_id = record.get('group_id')
+    if 'Entity_' + group_id.replace('-', '') in labels:
+        labels.remove('Entity_' + group_id.replace('-', ''))
+
     entity_node = EntityNode(
         uuid=record['uuid'],
         name=record['name'],
         name_embedding=record.get('name_embedding'),
-        group_id=record['group_id'],
-        labels=record['labels'],
+        group_id=group_id,
+        labels=labels,
         created_at=parse_db_date(record['created_at']),  # type: ignore
         summary=record['summary'],
-        attributes=record['attributes'],
+        attributes=attributes,
     )
 
-    entity_node.attributes.pop('uuid', None)
-    entity_node.attributes.pop('name', None)
-    entity_node.attributes.pop('group_id', None)
-    entity_node.attributes.pop('name_embedding', None)
-    entity_node.attributes.pop('summary', None)
-    entity_node.attributes.pop('created_at', None)
-
     return entity_node