diff --git a/graphiti_core/graphiti.py b/graphiti_core/graphiti.py index 0a80c4f6..3b0cf73c 100644 --- a/graphiti_core/graphiti.py +++ b/graphiti_core/graphiti.py @@ -380,6 +380,7 @@ class Graphiti: resolve_extracted_edges( self.clients, edges, + episode, ), extract_attributes_from_nodes( self.clients, nodes, episode, previous_episodes, entity_types @@ -682,7 +683,11 @@ class Graphiti: related_edges = await get_relevant_edges(self.driver, [updated_edge], SearchFilters(), 0.8) - resolved_edge = await dedupe_extracted_edge(self.llm_client, updated_edge, related_edges[0]) + resolved_edge = await dedupe_extracted_edge( + self.llm_client, + updated_edge, + related_edges[0], + ) contradicting_edges = await get_edge_contradictions(self.llm_client, edge, related_edges[0]) invalidated_edges = resolve_edge_contradictions(resolved_edge, contradicting_edges) diff --git a/graphiti_core/utils/maintenance/edge_operations.py b/graphiti_core/utils/maintenance/edge_operations.py index 2fd6e9b7..94cf010e 100644 --- a/graphiti_core/utils/maintenance/edge_operations.py +++ b/graphiti_core/utils/maintenance/edge_operations.py @@ -235,6 +235,7 @@ async def dedupe_extracted_edges( async def resolve_extracted_edges( clients: GraphitiClients, extracted_edges: list[EntityEdge], + episode: EpisodicNode, ) -> tuple[list[EntityEdge], list[EntityEdge]]: driver = clients.driver llm_client = clients.llm_client @@ -258,10 +259,7 @@ async def resolve_extracted_edges( await semaphore_gather( *[ resolve_extracted_edge( - llm_client, - extracted_edge, - related_edges, - existing_edges, + llm_client, extracted_edge, related_edges, existing_edges, episode ) for extracted_edge, related_edges, existing_edges in zip( extracted_edges, related_edges_lists, edge_invalidation_candidates, strict=True @@ -322,9 +320,10 @@ async def resolve_extracted_edge( extracted_edge: EntityEdge, related_edges: list[EntityEdge], existing_edges: list[EntityEdge], + episode: EpisodicNode, ) -> tuple[EntityEdge, list[EntityEdge]]: resolved_edge, invalidation_candidates = await semaphore_gather( - dedupe_extracted_edge(llm_client, extracted_edge, related_edges), + dedupe_extracted_edge(llm_client, extracted_edge, related_edges, episode), get_edge_contradictions(llm_client, extracted_edge, existing_edges), ) @@ -356,7 +355,10 @@ async def resolve_extracted_edge( async def dedupe_extracted_edge( - llm_client: LLMClient, extracted_edge: EntityEdge, related_edges: list[EntityEdge] + llm_client: LLMClient, + extracted_edge: EntityEdge, + related_edges: list[EntityEdge], + episode: EpisodicNode | None = None, ) -> EntityEdge: if len(related_edges) == 0: return extracted_edge @@ -391,6 +393,9 @@ async def dedupe_extracted_edge( else extracted_edge ) + if duplicate_fact_id >= 0 and episode is not None: + edge.episodes += episode.uuid + end = time() logger.debug( f'Resolved Edge: {extracted_edge.name} is {edge.name}, in {(end - start) * 1000} ms' diff --git a/graphiti_core/utils/maintenance/node_operations.py b/graphiti_core/utils/maintenance/node_operations.py index 9c27523e..8df4f563 100644 --- a/graphiti_core/utils/maintenance/node_operations.py +++ b/graphiti_core/utils/maintenance/node_operations.py @@ -437,7 +437,7 @@ async def dedupe_node_list( # Prepare context for LLM nodes_context = [ - {'uuid': node.uuid, 'name': node.name, 'summary': node.summary}.update(node.attributes) + {'uuid': node.uuid, 'name': node.name, 'summary': node.summary, **node.attributes} for node in nodes ]