graphiti-core 0.19.0rc3__py3-none-any.whl → 0.20.1__py3-none-any.whl

graphiti_core/graphiti.py

@@ -1070,7 +1070,7 @@ class Graphiti:
                 if record['episode_count'] == 1:
                     nodes_to_delete.append(node)
 
+        await Edge.delete_by_uuids(self.driver, [edge.uuid for edge in edges_to_delete])
         await Node.delete_by_uuids(self.driver, [node.uuid for node in nodes_to_delete])
 
-        await Edge.delete_by_uuids(self.driver, [edge.uuid for edge in edges_to_delete])
         await episode.delete(self.driver)

graphiti_core/helpers.py

@@ -26,7 +26,6 @@ from dotenv import load_dotenv
 from neo4j import time as neo4j_time
 from numpy._typing import NDArray
 from pydantic import BaseModel
-from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import GraphProvider
 from graphiti_core.errors import GroupIdValidationError
@@ -38,19 +37,15 @@ SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
 MAX_REFLEXION_ITERATIONS = int(os.getenv('MAX_REFLEXION_ITERATIONS', 0))
 DEFAULT_PAGE_LIMIT = 20
 
-RUNTIME_QUERY: LiteralString = (
-    'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
-)
 
+def parse_db_date(input_date: neo4j_time.DateTime | str | None) -> datetime | None:
+    if isinstance(input_date, neo4j_time.DateTime):
+        return input_date.to_native()
 
-def parse_db_date(neo_date: neo4j_time.DateTime | str | None) -> datetime | None:
-    return (
-        neo_date.to_native()
-        if isinstance(neo_date, neo4j_time.DateTime)
-        else datetime.fromisoformat(neo_date)
-        if neo_date
-        else None
-    )
+    if isinstance(input_date, str):
+        return datetime.fromisoformat(input_date)
+
+    return input_date
 
 
 def get_default_group_id(provider: GraphProvider) -> str:

graphiti_core/migrations/neo4j_node_group_labels.py

@@ -1,4 +1,5 @@
 import asyncio
+import csv
 import os
 
 from graphiti_core.driver.driver import GraphDriver
@@ -57,14 +58,41 @@ async def neo4j_node_group_labels(driver: GraphDriver, group_id: str, batch_size
         )
 
 
-async def neo4j_node_label_migration(driver: GraphDriver):
+def pop_last_n_group_ids(csv_file: str = 'group_ids.csv', count: int = 10):
+    with open(csv_file) as file:
+        reader = csv.reader(file)
+        group_ids = [row[0] for row in reader]
+
+    total_count = len(group_ids)
+    popped = group_ids[-count:]
+    remaining = group_ids[:-count]
+
+    with open(csv_file, 'w', newline='') as file:
+        writer = csv.writer(file)
+        for gid in remaining:
+            writer.writerow([gid])
+
+    return popped, total_count
+
+
+async def get_group_ids(driver: GraphDriver):
     query = """MATCH (n:Episodic)
                 RETURN DISTINCT n.group_id AS group_id"""
 
     results, _, _ = await driver.execute_query(query)
-    for result in results:
-        group_id = result['group_id']
-        await neo4j_node_group_labels(driver, group_id)
+    group_ids = [result['group_id'] for result in results]
+
+    with open('group_ids.csv', 'w', newline='') as file:
+        writer = csv.writer(file)
+        for gid in group_ids:
+            writer.writerow([gid])
+
+
+async def neo4j_node_label_migration(driver: GraphDriver, batch_size: int = 10):
+    group_ids, total = pop_last_n_group_ids(csv_file='group_ids.csv', count=batch_size)
+    while len(group_ids) > 0:
+        await asyncio.gather(*[neo4j_node_group_labels(driver, group_id) for group_id in group_ids])
+        group_ids, _ = pop_last_n_group_ids(csv_file='group_ids.csv', count=batch_size)
 
 
 async def main():
@@ -77,6 +105,7 @@ async def main():
         user=neo4j_user,
         password=neo4j_password,
     )
+    await get_group_ids(driver)
     await neo4j_node_label_migration(driver)
     await driver.close()
 

graphiti_core/models/edges/edge_db_queries.py

@@ -20,18 +20,36 @@ EPISODIC_EDGE_SAVE = """
     MATCH (episode:Episodic {uuid: $episode_uuid})
     MATCH (node:Entity {uuid: $entity_uuid})
     MERGE (episode)-[e:MENTIONS {uuid: $uuid}]->(node)
-    SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+    SET
+        e.group_id = $group_id,
+        e.created_at = $created_at
     RETURN e.uuid AS uuid
 """
 
-EPISODIC_EDGE_SAVE_BULK = """
-    UNWIND $episodic_edges AS edge
-    MATCH (episode:Episodic {uuid: edge.source_node_uuid})
-    MATCH (node:Entity {uuid: edge.target_node_uuid})
-    MERGE (episode)-[e:MENTIONS {uuid: edge.uuid}]->(node)
-    SET e = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at}
-    RETURN e.uuid AS uuid
-"""
+
+def get_episodic_edge_save_bulk_query(provider: GraphProvider) -> str:
+    if provider == GraphProvider.KUZU:
+        return """
+            MATCH (episode:Episodic {uuid: $source_node_uuid})
+            MATCH (node:Entity {uuid: $target_node_uuid})
+            MERGE (episode)-[e:MENTIONS {uuid: $uuid}]->(node)
+            SET
+                e.group_id = $group_id,
+                e.created_at = $created_at
+            RETURN e.uuid AS uuid
+        """
+
+    return """
+        UNWIND $episodic_edges AS edge
+        MATCH (episode:Episodic {uuid: edge.source_node_uuid})
+        MATCH (node:Entity {uuid: edge.target_node_uuid})
+        MERGE (episode)-[e:MENTIONS {uuid: edge.uuid}]->(node)
+        SET
+            e.group_id = edge.group_id,
+            e.created_at = edge.created_at
+        RETURN e.uuid AS uuid
+    """
+
 
 EPISODIC_EDGE_RETURN = """
     e.uuid AS uuid,
@@ -54,14 +72,32 @@ def get_entity_edge_save_query(provider: GraphProvider) -> str:
             """
         case GraphProvider.NEPTUNE:
             return """
-                MATCH (source:Entity {uuid: $edge_data.source_uuid}) 
-                MATCH (target:Entity {uuid: $edge_data.target_uuid}) 
+                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 = removeKeyFromMap(removeKeyFromMap($edge_data, "fact_embedding"), "episodes")
                 SET e.fact_embedding = join([x IN coalesce($edge_data.fact_embedding, []) | toString(x) ], ",")
                 SET e.episodes = join($edge_data.episodes, ",")
                 RETURN $edge_data.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MATCH (source:Entity {uuid: $source_uuid})
+                MATCH (target:Entity {uuid: $target_uuid})
+                MERGE (source)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: $uuid})-[:RELATES_TO]->(target)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at,
+                    e.name = $name,
+                    e.fact = $fact,
+                    e.fact_embedding = $fact_embedding,
+                    e.episodes = $episodes,
+                    e.expired_at = $expired_at,
+                    e.valid_at = $valid_at,
+                    e.invalid_at = $invalid_at,
+                    e.attributes = $attributes
+                RETURN e.uuid AS uuid
+            """
         case _:  # Neo4j
             return """
                 MATCH (source:Entity {uuid: $edge_data.source_uuid})
@@ -89,14 +125,32 @@ def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
         case GraphProvider.NEPTUNE:
             return """
                 UNWIND $entity_edges AS edge
-                MATCH (source:Entity {uuid: edge.source_node_uuid}) 
-                MATCH (target:Entity {uuid: edge.target_node_uuid}) 
+                MATCH (source:Entity {uuid: edge.source_node_uuid})
+                MATCH (target:Entity {uuid: edge.target_node_uuid})
                 MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
                 SET r = removeKeyFromMap(removeKeyFromMap(edge, "fact_embedding"), "episodes")
                 SET r.fact_embedding = join([x IN coalesce(edge.fact_embedding, []) | toString(x) ], ",")
                 SET r.episodes = join(edge.episodes, ",")
                 RETURN edge.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MATCH (source:Entity {uuid: $source_node_uuid})
+                MATCH (target:Entity {uuid: $target_node_uuid})
+                MERGE (source)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: $uuid})-[:RELATES_TO]->(target)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at,
+                    e.name = $name,
+                    e.fact = $fact,
+                    e.fact_embedding = $fact_embedding,
+                    e.episodes = $episodes,
+                    e.expired_at = $expired_at,
+                    e.valid_at = $valid_at,
+                    e.invalid_at = $invalid_at,
+                    e.attributes = $attributes
+                RETURN e.uuid AS uuid
+            """
         case _:
             return """
                 UNWIND $entity_edges AS edge
@@ -109,35 +163,42 @@ def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
             """
 
 
-ENTITY_EDGE_RETURN = """
-    e.uuid AS uuid,
-    n.uuid AS source_node_uuid,
-    m.uuid AS target_node_uuid,
-    e.group_id AS group_id,
-    e.name AS name,
-    e.fact AS fact,
-    e.episodes AS episodes,
-    e.created_at AS created_at,
-    e.expired_at AS expired_at,
-    e.valid_at AS valid_at,
-    e.invalid_at AS invalid_at,
-    properties(e) AS attributes
-"""
+def get_entity_edge_return_query(provider: GraphProvider) -> str:
+    # `fact_embedding` is not returned by default and must be manually loaded using `load_fact_embedding()`.
 
-ENTITY_EDGE_RETURN_NEPTUNE = """
-    e.uuid AS uuid,
-    n.uuid AS source_node_uuid,
-    m.uuid AS target_node_uuid,
-    e.group_id AS group_id,
-    e.name AS name,
-    e.fact AS fact,
-    split(e.episodes, ',') AS episodes,
-    e.created_at AS created_at,
-    e.expired_at AS expired_at,
-    e.valid_at AS valid_at,
-    e.invalid_at AS invalid_at,
-    properties(e) AS attributes
-"""
+    if provider == GraphProvider.NEPTUNE:
+        return """
+        e.uuid AS uuid,
+        n.uuid AS source_node_uuid,
+        m.uuid AS target_node_uuid,
+        e.group_id AS group_id,
+        e.name AS name,
+        e.fact AS fact,
+        split(e.episodes, ',') AS episodes,
+        e.created_at AS created_at,
+        e.expired_at AS expired_at,
+        e.valid_at AS valid_at,
+        e.invalid_at AS invalid_at,
+        properties(e) AS attributes
+    """
+
+    return """
+        e.uuid AS uuid,
+        n.uuid AS source_node_uuid,
+        m.uuid AS target_node_uuid,
+        e.group_id AS group_id,
+        e.created_at AS created_at,
+        e.name AS name,
+        e.fact AS fact,
+        e.episodes AS episodes,
+        e.expired_at AS expired_at,
+        e.valid_at AS valid_at,
+        e.invalid_at AS invalid_at,
+    """ + (
+        'e.attributes AS attributes'
+        if provider == GraphProvider.KUZU
+        else 'properties(e) AS attributes'
+    )
 
 
 def get_community_edge_save_query(provider: GraphProvider) -> str:
@@ -152,7 +213,7 @@ def get_community_edge_save_query(provider: GraphProvider) -> str:
             """
         case GraphProvider.NEPTUNE:
             return """
-                MATCH (community:Community {uuid: $community_uuid}) 
+                MATCH (community:Community {uuid: $community_uuid})
                 MATCH (node {uuid: $entity_uuid})
                 WHERE node:Entity OR node:Community
                 MERGE (community)-[r:HAS_MEMBER {uuid: $uuid}]->(node)
@@ -161,6 +222,24 @@ def get_community_edge_save_query(provider: GraphProvider) -> str:
                 SET r.created_at= $created_at
                 RETURN r.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node:Entity {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at
+                RETURN e.uuid AS uuid
+                UNION
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node:Community {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at
+                RETURN e.uuid AS uuid
+            """
         case _:  # Neo4j
             return """
                 MATCH (community:Community {uuid: $community_uuid})

graphiti_core/models/nodes/node_db_queries.py

@@ -24,10 +24,24 @@ def get_episode_node_save_query(provider: GraphProvider) -> str:
         case GraphProvider.NEPTUNE:
             return """
                 MERGE (n:Episodic {uuid: $uuid})
-                SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content, 
+                SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
                 entity_edges: join([x IN coalesce($entity_edges, []) | toString(x) ], '|'), created_at: $created_at, valid_at: $valid_at}
                 RETURN n.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MERGE (n:Episodic {uuid: $uuid})
+                SET
+                    n.name = $name,
+                    n.group_id = $group_id,
+                    n.created_at = $created_at,
+                    n.source = $source,
+                    n.source_description = $source_description,
+                    n.content = $content,
+                    n.valid_at = $valid_at,
+                    n.entity_edges = $entity_edges
+                RETURN n.uuid AS uuid
+            """
         case GraphProvider.FALKORDB:
             return """
                 MERGE (n:Episodic {uuid: $uuid})
@@ -51,11 +65,25 @@ def get_episode_node_save_bulk_query(provider: GraphProvider) -> str:
             return """
                 UNWIND $episodes AS episode
                 MERGE (n:Episodic {uuid: episode.uuid})
-                SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description, 
-                    source: episode.source, content: episode.content, 
+                SET n = {uuid: episode.uuid, name: episode.name, group_id: episode.group_id, source_description: episode.source_description,
+                    source: episode.source, content: episode.content,
                 entity_edges: join([x IN coalesce(episode.entity_edges, []) | toString(x) ], '|'), created_at: episode.created_at, valid_at: episode.valid_at}
                 RETURN n.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MERGE (n:Episodic {uuid: $uuid})
+                SET
+                    n.name = $name,
+                    n.group_id = $group_id,
+                    n.created_at = $created_at,
+                    n.source = $source,
+                    n.source_description = $source_description,
+                    n.content = $content,
+                    n.valid_at = $valid_at,
+                    n.entity_edges = $entity_edges
+                RETURN n.uuid AS uuid
+            """
         case GraphProvider.FALKORDB:
             return """
                 UNWIND $episodes AS episode
@@ -76,14 +104,14 @@ def get_episode_node_save_bulk_query(provider: GraphProvider) -> str:
 
 
 EPISODIC_NODE_RETURN = """
-    e.content AS content,
-    e.created_at AS created_at,
-    e.valid_at AS valid_at,
     e.uuid AS uuid,
     e.name AS name,
     e.group_id AS group_id,
-    e.source_description AS source_description,
+    e.created_at AS created_at,
     e.source AS source,
+    e.source_description AS source_description,
+    e.content AS content,
+    e.valid_at AS valid_at,
     e.entity_edges AS entity_edges
 """
 
@@ -109,6 +137,20 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
                 SET n = $entity_data
                 RETURN n.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MERGE (n:Entity {uuid: $uuid})
+                SET
+                    n.name = $name,
+                    n.group_id = $group_id,
+                    n.labels = $labels,
+                    n.created_at = $created_at,
+                    n.name_embedding = $name_embedding,
+                    n.summary = $summary,
+                    n.attributes = $attributes
+                WITH n
+                RETURN n.uuid AS uuid
+            """
         case GraphProvider.NEPTUNE:
             label_subquery = ''
             for label in labels.split(':'):
@@ -168,6 +210,19 @@ def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict])
                     """
                 )
             return queries
+        case GraphProvider.KUZU:
+            return """
+                MERGE (n:Entity {uuid: $uuid})
+                SET
+                    n.name = $name,
+                    n.group_id = $group_id,
+                    n.labels = $labels,
+                    n.created_at = $created_at,
+                    n.name_embedding = $name_embedding,
+                    n.summary = $summary,
+                    n.attributes = $attributes
+                RETURN n.uuid AS uuid
+            """
         case _:  # Neo4j
             return """
                 UNWIND $nodes AS node
@@ -179,15 +234,28 @@ def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict])
             """
 
 
-ENTITY_NODE_RETURN = """
-    n.uuid AS uuid,
-    n.name AS name,
-    n.group_id AS group_id,
-    n.created_at AS created_at,
-    n.summary AS summary,
-    labels(n) AS labels,
-    properties(n) AS attributes
-"""
+def get_entity_node_return_query(provider: GraphProvider) -> str:
+    # `name_embedding` is not returned by default and must be loaded manually using `load_name_embedding()`.
+    if provider == GraphProvider.KUZU:
+        return """
+            n.uuid AS uuid,
+            n.name AS name,
+            n.group_id AS group_id,
+            n.labels AS labels,
+            n.created_at AS created_at,
+            n.summary AS summary,
+            n.attributes AS attributes
+        """
+
+    return """
+        n.uuid AS uuid,
+        n.name AS name,
+        n.group_id AS group_id,
+        n.created_at AS created_at,
+        n.summary AS summary,
+        labels(n) AS labels,
+        properties(n) AS attributes
+    """
 
 
 def get_community_node_save_query(provider: GraphProvider) -> str:
@@ -201,10 +269,21 @@ def get_community_node_save_query(provider: GraphProvider) -> str:
         case GraphProvider.NEPTUNE:
             return """
                 MERGE (n:Community {uuid: $uuid})
-                SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}        
+                SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
                 SET n.name_embedding = join([x IN coalesce($name_embedding, []) | toString(x) ], ",")
                 RETURN n.uuid AS uuid
             """
+        case GraphProvider.KUZU:
+            return """
+                MERGE (n:Community {uuid: $uuid})
+                SET
+                    n.name = $name,
+                    n.group_id = $group_id,
+                    n.created_at = $created_at,
+                    n.name_embedding = $name_embedding,
+                    n.summary = $summary
+                RETURN n.uuid AS uuid
+            """
         case _:  # Neo4j
             return """
                 MERGE (n:Community {uuid: $uuid})
@@ -215,12 +294,12 @@ def get_community_node_save_query(provider: GraphProvider) -> str:
 
 
 COMMUNITY_NODE_RETURN = """
-    n.uuid AS uuid,
-    n.name AS name,
-    n.name_embedding AS name_embedding,
-    n.group_id AS group_id,
-    n.summary AS summary,
-    n.created_at AS created_at
+    c.uuid AS uuid,
+    c.name AS name,
+    c.group_id AS group_id,
+    c.created_at AS created_at,
+    c.name_embedding AS name_embedding,
+    c.summary AS summary
 """
 
 COMMUNITY_NODE_RETURN_NEPTUNE = """