working refactor

prasmussen15 · prasmussen15 · commit 51bf29fe0529 · 2025-04-25T23:53:24.000-04:00
diff --git a/graphiti_core/models/edges/edge_db_queries.py b/graphiti_core/models/edges/edge_db_queries.py
@@ -23,8 +23,8 @@
 
 EPISODIC_EDGE_SAVE_BULK = """
     UNWIND $episodic_edges AS edge
-    MATCH (episode:Episodic {uuid: startNode(edge).uuid}) 
-    MATCH (node:Entity {uuid: endNode(edge).uuid}) 
+    MATCH (episode:Episodic {uuid: edge.source_node_uuid}) 
+    MATCH (node:Entity {uuid: edge.target_node_uuid}) 
     MERGE (episode)-[r:MENTIONS {uuid: edge.uuid}]->(node)
     SET r = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at}
     RETURN r.uuid AS uuid
@@ -47,7 +47,7 @@
     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}
     WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
-    RETURN r.uuid AS uuid
+    RETURN edge.uuid AS uuid
 """
 
 COMMUNITY_EDGE_SAVE = """
diff --git a/graphiti_core/search/search_utils.py b/graphiti_core/search/search_utils.py
@@ -341,10 +341,10 @@ async def node_fulltext_search(
 
     query = (
         """
-                                        CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
-                                        YIELD node AS n, score
-                                        WHERE n:Entity
-                                        """
+                                            CALL db.index.fulltext.queryNodes("node_name_and_summary", $query, {limit: $limit}) 
+                                            YIELD node AS n, score
+                                            WHERE n:Entity
+                                            """
         + filter_query
         + ENTITY_NODE_RETURN
         + """
@@ -739,12 +739,25 @@ async def get_relevant_edges(
             """
         + filter_query
         + """
-            WITH n, m, e, vector.similarity.cosine(e.fact_embedding, edge.fact_embedding) AS score
-            WHERE score > $min_score"""
-        + ENTITY_EDGE_RETURN
-        + """
-        ORDER BY score DESC
-        LIMIT $limit
+            WITH e, edge, vector.similarity.cosine(e.fact_embedding, edge.fact_embedding) AS score
+            WHERE score > $min_score
+            WITH edge, e, score
+            ORDER BY score DESC
+            RETURN edge.uuid AS search_edge_uuid,
+                collect({
+                    uuid: e.uuid,
+                    source_node_uuid: startNode(e).uuid,
+                    target_node_uuid: endNode(e).uuid,
+                    created_at: e.created_at,
+                    name: e.name,
+                    group_id: e.group_id,
+                    fact: e.fact,
+                    fact_embedding: e.fact_embedding,
+                    episodes: e.episodes,
+                    expired_at: e.expired_at,
+                    valid_at: e.valid_at,
+                    invalid_at: e.invalid_at
+                })[..$limit] AS matches
         """
     )
 
@@ -792,15 +805,29 @@ async def get_edge_invalidation_candidates(
             """
         + filter_query
         + """
-            WITH e, vector.similarity.cosine(e.fact_embedding, edge.fact_embedding) AS score
+            WITH edge, e, vector.similarity.cosine(e.fact_embedding, edge.fact_embedding) AS score
             WHERE score > $min_score
-            
-        ORDER BY score DESC
-        LIMIT $limit
+            WITH edge, e, score
+            ORDER BY score DESC
+            RETURN edge.uuid AS search_edge_uuid,
+                collect({
+                    uuid: e.uuid,
+                    source_node_uuid: startNode(e).uuid,
+                    target_node_uuid: endNode(e).uuid,
+                    created_at: e.created_at,
+                    name: e.name,
+                    group_id: e.group_id,
+                    fact: e.fact,
+                    fact_embedding: e.fact_embedding,
+                    episodes: e.episodes,
+                    expired_at: e.expired_at,
+                    valid_at: e.valid_at,
+                    invalid_at: e.invalid_at
+                })[..$limit] AS matches
         """
     )
 
-    records_list, _, _ = await driver.execute_query(
+    results, _, _ = await driver.execute_query(
         query,
         query_params,
         edges=[edge.model_dump() for edge in edges],
@@ -809,11 +836,16 @@ async def get_edge_invalidation_candidates(
         database_=DEFAULT_DATABASE,
         routing_='r',
     )
-    relevant_edges = [
-        [get_entity_edge_from_record(record) for record in records] for records in records_list
-    ]
+    invalidation_edges_dict: dict[str, list[EntityEdge]] = {
+        result['search_edge_uuid']: [
+            get_entity_edge_from_record(record) for record in result['matches']
+        ]
+        for result in results
+    }
 
-    return relevant_edges
+    invalidation_edges = [invalidation_edges_dict.get(edge.uuid, []) for edge in edges]
+
+    return invalidation_edges
 
 
 # takes in a list of rankings of uuids
diff --git a/graphiti_core/utils/bulk_utils.py b/graphiti_core/utils/bulk_utils.py
@@ -71,7 +71,7 @@ class RawEpisode(BaseModel):
 
 
 async def retrieve_previous_episodes_bulk(
-    driver: AsyncDriver, episodes: list[EpisodicNode]
+        driver: AsyncDriver, episodes: list[EpisodicNode]
 ) -> list[tuple[EpisodicNode, list[EpisodicNode]]]:
     previous_episodes_list = await semaphore_gather(
         *[
@@ -89,11 +89,11 @@ async def retrieve_previous_episodes_bulk(
 
 
 async def add_nodes_and_edges_bulk(
-    driver: AsyncDriver,
-    episodic_nodes: list[EpisodicNode],
-    episodic_edges: list[EpisodicEdge],
-    entity_nodes: list[EntityNode],
-    entity_edges: list[EntityEdge],
+        driver: AsyncDriver,
+        episodic_nodes: list[EpisodicNode],
+        episodic_edges: list[EpisodicEdge],
+        entity_nodes: list[EntityNode],
+        entity_edges: list[EntityEdge],
 ):
     async with driver.session(database=DEFAULT_DATABASE) as session:
         await session.execute_write(
@@ -102,11 +102,11 @@ async def add_nodes_and_edges_bulk(
 
 
 async def add_nodes_and_edges_bulk_tx(
-    tx: AsyncManagedTransaction,
-    episodic_nodes: list[EpisodicNode],
-    episodic_edges: list[EpisodicEdge],
-    entity_nodes: list[EntityNode],
-    entity_edges: list[EntityEdge],
+        tx: AsyncManagedTransaction,
+        episodic_nodes: list[EpisodicNode],
+        episodic_edges: list[EpisodicEdge],
+        entity_nodes: list[EntityNode],
+        entity_edges: list[EntityEdge],
 ):
     episodes = [dict(episode) for episode in episodic_nodes]
     for episode in episodes:
@@ -128,12 +128,14 @@ async def add_nodes_and_edges_bulk_tx(
 
     await tx.run(EPISODIC_NODE_SAVE_BULK, episodes=episodes)
     await tx.run(ENTITY_NODE_SAVE_BULK, nodes=nodes)
-    await tx.run(EPISODIC_EDGE_SAVE_BULK, episodic_edges=[dict(edge) for edge in episodic_edges])
-    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=[dict(edge) for edge in entity_edges])
+    await tx.run(
+        EPISODIC_EDGE_SAVE_BULK, episodic_edges=[edge.model_dump() for edge in episodic_edges]
+    )
+    await tx.run(ENTITY_EDGE_SAVE_BULK, entity_edges=[edge.model_dump() for edge in entity_edges])
 
 
 async def extract_nodes_and_edges_bulk(
-    llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
+        llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
 ) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:
     extracted_nodes_bulk = await semaphore_gather(
         *[
@@ -176,16 +178,16 @@ async def extract_nodes_and_edges_bulk(
 
 
 async def dedupe_nodes_bulk(
-    driver: AsyncDriver,
-    llm_client: LLMClient,
-    extracted_nodes: list[EntityNode],
+        driver: AsyncDriver,
+        llm_client: LLMClient,
+        extracted_nodes: list[EntityNode],
 ) -> tuple[list[EntityNode], dict[str, str]]:
     # Compress nodes
     nodes, uuid_map = node_name_match(extracted_nodes)
 
     compressed_nodes, compressed_map = await compress_nodes(llm_client, nodes, uuid_map)
 
-    node_chunks = [nodes[i : i + CHUNK_SIZE] for i in range(0, len(nodes), CHUNK_SIZE)]
+    node_chunks = [nodes[i: i + CHUNK_SIZE] for i in range(0, len(nodes), CHUNK_SIZE)]
 
     existing_nodes_chunks: list[list[EntityNode]] = list(
         await semaphore_gather(
@@ -212,13 +214,13 @@ async def dedupe_nodes_bulk(
 
 
 async def dedupe_edges_bulk(
-    driver: AsyncDriver, llm_client: LLMClient, extracted_edges: list[EntityEdge]
+        driver: AsyncDriver, llm_client: LLMClient, extracted_edges: list[EntityEdge]
 ) -> list[EntityEdge]:
     # First compress edges
     compressed_edges = await compress_edges(llm_client, extracted_edges)
 
     edge_chunks = [
-        compressed_edges[i : i + CHUNK_SIZE] for i in range(0, len(compressed_edges), CHUNK_SIZE)
+        compressed_edges[i: i + CHUNK_SIZE] for i in range(0, len(compressed_edges), CHUNK_SIZE)
     ]
 
     relevant_edges_chunks: list[list[EntityEdge]] = list(
@@ -254,7 +256,7 @@ def node_name_match(nodes: list[EntityNode]) -> tuple[list[EntityNode], dict[str
 
 
 async def compress_nodes(
-    llm_client: LLMClient, nodes: list[EntityNode], uuid_map: dict[str, str]
+        llm_client: LLMClient, nodes: list[EntityNode], uuid_map: dict[str, str]
 ) -> tuple[list[EntityNode], dict[str, str]]:
     # We want to first compress the nodes by deduplicating nodes across each of the episodes added in bulk
     if len(nodes) == 0:
@@ -375,9 +377,9 @@ def resolve_edge_pointers(edges: list[E], uuid_map: dict[str, str]):
 
 
 async def extract_edge_dates_bulk(
-    llm_client: LLMClient,
-    extracted_edges: list[EntityEdge],
-    episode_pairs: list[tuple[EpisodicNode, list[EpisodicNode]]],
+        llm_client: LLMClient,
+        extracted_edges: list[EntityEdge],
+        episode_pairs: list[tuple[EpisodicNode, list[EpisodicNode]]],
 ) -> list[EntityEdge]:
     edges: list[EntityEdge] = []
     # confirm that all of our edges have at least one episode
