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

graphiti_core/migrations/neo4j_node_group_labels.py

@@ -0,0 +1,114 @@
+import asyncio
+import csv
+import os
+
+from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.driver.neo4j_driver import Neo4jDriver
+from graphiti_core.helpers import validate_group_id
+from graphiti_core.utils.maintenance.graph_data_operations import build_dynamic_indexes
+
+
+async def neo4j_node_group_labels(driver: GraphDriver, group_id: str, batch_size: int = 100):
+    validate_group_id(group_id)
+    await build_dynamic_indexes(driver, group_id)
+
+    episode_query = """
+                        MATCH (n:Episodic {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            SET n:$($group_label)
+                        } IN TRANSACTIONS OF $batch_size ROWS"""
+
+    entity_query = """
+                        MATCH (n:Entity {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            SET n:$($group_label)
+                        } IN TRANSACTIONS OF $batch_size ROWS"""
+
+    community_query = """
+                        MATCH (n:Community {group_id: $group_id})
+                        CALL {
+                            WITH n
+                            SET n:$($group_label)
+                        } IN TRANSACTIONS OF $batch_size ROWS"""
+
+    async with driver.session() as session:
+        await session.run(
+            episode_query,
+            group_id=group_id,
+            group_label='Episodic_' + group_id.replace('-', ''),
+            batch_size=batch_size,
+        )
+
+    async with driver.session() as session:
+        await session.run(
+            entity_query,
+            group_id=group_id,
+            group_label='Entity_' + group_id.replace('-', ''),
+            batch_size=batch_size,
+        )
+
+    async with driver.session() as session:
+        await session.run(
+            community_query,
+            group_id=group_id,
+            group_label='Community_' + group_id.replace('-', ''),
+            batch_size=batch_size,
+        )
+
+
+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)
+    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():
+    neo4j_uri = os.environ.get('NEO4J_URI') or 'bolt://localhost:7687'
+    neo4j_user = os.environ.get('NEO4J_USER') or 'neo4j'
+    neo4j_password = os.environ.get('NEO4J_PASSWORD') or 'password'
+
+    driver = Neo4jDriver(
+        uri=neo4j_uri,
+        user=neo4j_user,
+        password=neo4j_password,
+    )
+    await get_group_ids(driver)
+    await neo4j_node_label_migration(driver)
+    await driver.close()
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

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,
@@ -43,82 +61,193 @@ EPISODIC_EDGE_RETURN = """
 
 
 def get_entity_edge_save_query(provider: GraphProvider) -> str:
-    if provider == GraphProvider.FALKORDB:
-        return """
-            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 = $edge_data
-            RETURN e.uuid AS uuid
-        """
-
-    return """
-        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 = $edge_data
-        WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)
-        RETURN e.uuid AS uuid
-    """
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                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 = $edge_data
+                RETURN e.uuid AS uuid
+            """
+        case GraphProvider.NEPTUNE:
+            return """
+                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})
+                MATCH (target:Entity {uuid: $edge_data.target_uuid})
+                MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+                SET e = $edge_data
+                WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)
+                RETURN e.uuid AS uuid
+            """
 
 
 def get_entity_edge_save_bulk_query(provider: GraphProvider) -> str:
-    if provider == GraphProvider.FALKORDB:
-        return """
-            UNWIND $entity_edges AS edge
-            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 = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes,
-            created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at, fact_embedding: vecf32(edge.fact_embedding)}
-            WITH r, edge
-            RETURN edge.uuid AS uuid
-        """
-
-    return """
-        UNWIND $entity_edges AS edge
-        MATCH (source:Entity {uuid: edge.source_node_uuid})
-        MATCH (target:Entity {uuid: edge.target_node_uuid})
-        MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
-        SET e = edge
-        WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)
-        RETURN edge.uuid AS uuid
-    """
-
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                UNWIND $entity_edges AS edge
+                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 = {uuid: edge.uuid, name: edge.name, group_id: edge.group_id, fact: edge.fact, episodes: edge.episodes,
+                created_at: edge.created_at, expired_at: edge.expired_at, valid_at: edge.valid_at, invalid_at: edge.invalid_at, fact_embedding: vecf32(edge.fact_embedding)}
+                WITH r, edge
+                RETURN edge.uuid AS uuid
+            """
+        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})
+                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
+                MATCH (source:Entity {uuid: edge.source_node_uuid})
+                MATCH (target:Entity {uuid: edge.target_node_uuid})
+                MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
+                SET e = edge
+                WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)
+                RETURN edge.uuid AS uuid
+            """
 
-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()`.
 
-def get_community_edge_save_query(provider: GraphProvider) -> str:
-    if provider == GraphProvider.FALKORDB:
+    if provider == GraphProvider.NEPTUNE:
         return """
-            MATCH (community:Community {uuid: $community_uuid})
-            MATCH (node {uuid: $entity_uuid})
-            MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
-            SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
-            RETURN e.uuid AS uuid
-        """
+        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 """
-        MATCH (community:Community {uuid: $community_uuid})
-        MATCH (node:Entity | Community {uuid: $entity_uuid})
-        MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
-        SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
-        RETURN e.uuid AS uuid
-    """
+        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:
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+                RETURN e.uuid AS uuid
+            """
+        case GraphProvider.NEPTUNE:
+            return """
+                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)
+                SET r.uuid= $uuid
+                SET r.group_id= $group_id
+                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})
+                MATCH (node:Entity | Community {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+                RETURN e.uuid AS uuid
+            """
 
 
 COMMUNITY_EDGE_RETURN = """