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

graphiti_core/models/edges/edge_db_queries.py

@@ -43,47 +43,70 @@ 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 _:  # 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 _:
+            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 = """
@@ -101,24 +124,51 @@ ENTITY_EDGE_RETURN = """
     properties(e) AS attributes
 """
 
+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
+"""
+
 
 def get_community_edge_save_query(provider: GraphProvider) -> str:
-    if provider == 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
-        """
-
-    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
-    """
+    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 _:  # 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 = """

graphiti_core/models/nodes/node_db_queries.py

@@ -18,21 +18,62 @@ from typing import Any
 
 from graphiti_core.driver.driver import GraphProvider
 
-EPISODIC_NODE_SAVE = """
-    MERGE (n:Episodic {uuid: $uuid})
-    SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
-    entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
-    RETURN n.uuid AS uuid
-"""
 
-EPISODIC_NODE_SAVE_BULK = """
-    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,
-    entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
-    RETURN n.uuid AS uuid
-"""
+def get_episode_node_save_query(provider: GraphProvider) -> str:
+    match provider:
+        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, 
+                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.FALKORDB:
+            return """
+                MERGE (n:Episodic {uuid: $uuid})
+                SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
+                entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
+                RETURN n.uuid AS uuid
+            """
+        case _:  # Neo4j
+            return """
+                MERGE (n:Episodic {uuid: $uuid})
+                SET n:$($group_label)
+                SET n = {uuid: $uuid, name: $name, group_id: $group_id, source_description: $source_description, source: $source, content: $content,
+                entity_edges: $entity_edges, created_at: $created_at, valid_at: $valid_at}
+                RETURN n.uuid AS uuid
+            """
+
+
+def get_episode_node_save_bulk_query(provider: GraphProvider) -> str:
+    match provider:
+        case GraphProvider.NEPTUNE:
+            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, 
+                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.FALKORDB:
+            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, 
+                entity_edges: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
+                RETURN n.uuid AS uuid
+            """
+        case _:  # Neo4j
+            return """
+                UNWIND $episodes AS episode
+                MERGE (n:Episodic {uuid: episode.uuid})
+                SET n:$(episode.group_label)
+                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: episode.entity_edges, created_at: episode.created_at, valid_at: episode.valid_at}
+                RETURN n.uuid AS uuid
+            """
+
 
 EPISODIC_NODE_RETURN = """
     e.content AS content,
@@ -46,54 +87,96 @@ EPISODIC_NODE_RETURN = """
     e.entity_edges AS entity_edges
 """
 
+EPISODIC_NODE_RETURN_NEPTUNE = """
+    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.source AS source,
+    split(e.entity_edges, ",") AS entity_edges
+"""
+
 
 def get_entity_node_save_query(provider: GraphProvider, labels: str) -> str:
-    if provider == GraphProvider.FALKORDB:
-        return f"""
-            MERGE (n:Entity {{uuid: $entity_data.uuid}})
-            SET n:{labels}
-            SET n = $entity_data
-            RETURN n.uuid AS uuid
-        """
-
-    return f"""
-        MERGE (n:Entity {{uuid: $entity_data.uuid}})
-        SET n:{labels}
-        SET n = $entity_data
-        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)
-        RETURN n.uuid AS uuid
-    """
+    match provider:
+        case GraphProvider.FALKORDB:
+            return f"""
+                MERGE (n:Entity {{uuid: $entity_data.uuid}})
+                SET n:{labels}
+                SET n = $entity_data
+                RETURN n.uuid AS uuid
+            """
+        case GraphProvider.NEPTUNE:
+            label_subquery = ''
+            for label in labels.split(':'):
+                label_subquery += f' SET n:{label}\n'
+            return f"""
+                MERGE (n:Entity {{uuid: $entity_data.uuid}})
+                {label_subquery}
+                SET n = removeKeyFromMap(removeKeyFromMap($entity_data, "labels"), "name_embedding")
+                SET n.name_embedding = join([x IN coalesce($entity_data.name_embedding, []) | toString(x) ], ",")
+                RETURN n.uuid AS uuid
+            """
+        case _:
+            return f"""
+                MERGE (n:Entity {{uuid: $entity_data.uuid}})
+                SET n:{labels}
+                SET n = $entity_data
+                WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $entity_data.name_embedding)
+                RETURN n.uuid AS uuid
+            """
 
 
 def get_entity_node_save_bulk_query(provider: GraphProvider, nodes: list[dict]) -> str | Any:
-    if provider == GraphProvider.FALKORDB:
-        queries = []
-        for node in nodes:
-            for label in node['labels']:
+    match provider:
+        case GraphProvider.FALKORDB:
+            queries = []
+            for node in nodes:
+                for label in node['labels']:
+                    queries.append(
+                        (
+                            f"""
+                            UNWIND $nodes AS node
+                            MERGE (n:Entity {{uuid: node.uuid}})
+                            SET n:{label}
+                            SET n = node
+                            WITH n, node
+                            SET n.name_embedding = vecf32(node.name_embedding)
+                            RETURN n.uuid AS uuid
+                            """,
+                            {'nodes': [node]},
+                        )
+                    )
+            return queries
+        case GraphProvider.NEPTUNE:
+            queries = []
+            for node in nodes:
+                labels = ''
+                for label in node['labels']:
+                    labels += f' SET n:{label}\n'
                 queries.append(
-                    (
-                        f"""
+                    f"""
                         UNWIND $nodes AS node
                         MERGE (n:Entity {{uuid: node.uuid}})
-                        SET n:{label}
-                        SET n = node
-                        WITH n, node
-                        SET n.name_embedding = vecf32(node.name_embedding)
+                        {labels}
+                        SET n = removeKeyFromMap(removeKeyFromMap(node, "labels"), "name_embedding")
+                        SET n.name_embedding = join([x IN coalesce(node.name_embedding, []) | toString(x) ], ",")
                         RETURN n.uuid AS uuid
-                        """,
-                        {'nodes': [node]},
-                    )
+                    """
                 )
-        return queries
-
-    return """
-        UNWIND $nodes AS node
-        MERGE (n:Entity {uuid: node.uuid})
-        SET n:$(node.labels)
-        SET n = node
-        WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
-        RETURN n.uuid AS uuid
-    """
+            return queries
+        case _:  # Neo4j
+            return """
+                UNWIND $nodes AS node
+                MERGE (n:Entity {uuid: node.uuid})
+                SET n:$(node.labels)
+                SET n = node
+                WITH n, node CALL db.create.setNodeVectorProperty(n, "name_embedding", node.name_embedding)
+                RETURN n.uuid AS uuid
+            """
 
 
 ENTITY_NODE_RETURN = """
@@ -108,19 +191,27 @@ ENTITY_NODE_RETURN = """
 
 
 def get_community_node_save_query(provider: GraphProvider) -> str:
-    if provider == GraphProvider.FALKORDB:
-        return """
-            MERGE (n:Community {uuid: $uuid})
-            SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at, name_embedding: vecf32($name_embedding)}
-            RETURN n.uuid AS uuid
-        """
-
-    return """
-        MERGE (n:Community {uuid: $uuid})
-        SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
-        WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
-        RETURN n.uuid AS uuid
-    """
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                MERGE (n:Community {uuid: $uuid})
+                SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at, name_embedding: vecf32($name_embedding)}
+                RETURN n.uuid AS uuid
+            """
+        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.name_embedding = join([x IN coalesce($name_embedding, []) | toString(x) ], ",")
+                RETURN n.uuid AS uuid
+            """
+        case _:  # Neo4j
+            return """
+                MERGE (n:Community {uuid: $uuid})
+                SET n = {uuid: $uuid, name: $name, group_id: $group_id, summary: $summary, created_at: $created_at}
+                WITH n CALL db.create.setNodeVectorProperty(n, "name_embedding", $name_embedding)
+                RETURN n.uuid AS uuid
+            """
 
 
 COMMUNITY_NODE_RETURN = """
@@ -131,3 +222,12 @@ COMMUNITY_NODE_RETURN = """
     n.summary AS summary,
     n.created_at AS created_at
 """
+
+COMMUNITY_NODE_RETURN_NEPTUNE = """
+    n.uuid AS uuid,
+    n.name AS name,
+    [x IN split(n.name_embedding, ",") | toFloat(x)] AS name_embedding,
+    n.group_id AS group_id,
+    n.summary AS summary,
+    n.created_at AS created_at
+"""