""" Copyright 2024, Zep Software, Inc. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ import json import logging import typing from datetime import datetime import numpy as np from pydantic import BaseModel, Field from typing_extensions import Any from graphiti_core.driver.driver import ( GraphDriver, GraphDriverSession, GraphProvider, ) from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge, create_entity_edge_embeddings from graphiti_core.embedder import EmbedderClient from graphiti_core.graphiti_types import GraphitiClients from graphiti_core.helpers import normalize_l2, semaphore_gather from graphiti_core.models.edges.edge_db_queries import ( get_entity_edge_save_bulk_query, get_episodic_edge_save_bulk_query, ) from graphiti_core.models.nodes.node_db_queries import ( get_entity_node_save_bulk_query, get_episode_node_save_bulk_query, ) from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings from graphiti_core.utils.maintenance.dedup_helpers import ( DedupResolutionState, _build_candidate_indexes, _normalize_string_exact, _resolve_with_similarity, ) from graphiti_core.utils.maintenance.edge_operations import ( extract_edges, resolve_extracted_edge, ) from graphiti_core.utils.maintenance.graph_data_operations import ( EPISODE_WINDOW_LEN, retrieve_episodes, ) from graphiti_core.utils.maintenance.node_operations import ( extract_nodes, resolve_extracted_nodes, ) logger = logging.getLogger(__name__) CHUNK_SIZE = 10 def _build_directed_uuid_map(pairs: list[tuple[str, str]]) -> dict[str, str]: """Collapse alias -> canonical chains while preserving direction. The incoming pairs represent directed mappings discovered during node dedupe. We use a simple union-find with iterative path compression to ensure every source UUID resolves to its ultimate canonical target, even if aliases appear lexicographically smaller than the canonical UUID. """ parent: dict[str, str] = {} def find(uuid: str) -> str: """Directed union-find lookup using iterative path compression.""" parent.setdefault(uuid, uuid) root = uuid while parent[root] != root: root = parent[root] while parent[uuid] != root: next_uuid = parent[uuid] parent[uuid] = root uuid = next_uuid return root for source_uuid, target_uuid in pairs: parent.setdefault(source_uuid, source_uuid) parent.setdefault(target_uuid, target_uuid) parent[find(source_uuid)] = find(target_uuid) return {uuid: find(uuid) for uuid in parent} class RawEpisode(BaseModel): name: str uuid: str | None = Field(default=None) content: str source_description: str source: EpisodeType reference_time: datetime async def retrieve_previous_episodes_bulk( driver: GraphDriver, episodes: list[EpisodicNode] ) -> list[tuple[EpisodicNode, list[EpisodicNode]]]: previous_episodes_list = await semaphore_gather( *[ retrieve_episodes( driver, episode.valid_at, last_n=EPISODE_WINDOW_LEN, group_ids=[episode.group_id] ) for episode in episodes ] ) episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]] = [ (episode, previous_episodes_list[i]) for i, episode in enumerate(episodes) ] return episode_tuples async def add_nodes_and_edges_bulk( driver: GraphDriver, episodic_nodes: list[EpisodicNode], episodic_edges: list[EpisodicEdge], entity_nodes: list[EntityNode], entity_edges: list[EntityEdge], embedder: EmbedderClient, ): session = driver.session() try: await session.execute_write( add_nodes_and_edges_bulk_tx, episodic_nodes, episodic_edges, entity_nodes, entity_edges, embedder, driver=driver, ) finally: await session.close() async def add_nodes_and_edges_bulk_tx( tx: GraphDriverSession, episodic_nodes: list[EpisodicNode], episodic_edges: list[EpisodicEdge], entity_nodes: list[EntityNode], entity_edges: list[EntityEdge], embedder: EmbedderClient, driver: GraphDriver, ): episodes = [dict(episode) for episode in episodic_nodes] for episode in episodes: episode['source'] = str(episode['source'].value) episode.pop('labels', None) nodes = [] for node in entity_nodes: if node.name_embedding is None: await node.generate_name_embedding(embedder) entity_data: dict[str, Any] = { 'uuid': node.uuid, 'name': node.name, 'group_id': node.group_id, 'summary': node.summary, 'created_at': node.created_at, 'name_embedding': node.name_embedding, 'labels': list(set(node.labels + ['Entity'])), } if driver.provider == GraphProvider.KUZU: attributes = convert_datetimes_to_strings(node.attributes) if node.attributes else {} entity_data['attributes'] = json.dumps(attributes) else: for k, v in (node.attributes or {}).items(): if k not in entity_data: entity_data[k] = v nodes.append(entity_data) edges = [] for edge in entity_edges: if edge.fact_embedding is None: await edge.generate_embedding(embedder) edge_data: dict[str, Any] = { 'uuid': edge.uuid, 'source_node_uuid': edge.source_node_uuid, 'target_node_uuid': edge.target_node_uuid, 'name': edge.name, 'fact': edge.fact, 'group_id': edge.group_id, 'episodes': edge.episodes, 'created_at': edge.created_at, 'expired_at': edge.expired_at, 'valid_at': edge.valid_at, 'invalid_at': edge.invalid_at, 'reference_time': edge.reference_time, 'fact_embedding': edge.fact_embedding, } if driver.provider == GraphProvider.KUZU: attributes = convert_datetimes_to_strings(edge.attributes) if edge.attributes else {} edge_data['attributes'] = json.dumps(attributes) else: # Merge attributes without overwriting explicit edge fields. # Attributes may contain stale string versions of typed fields # (e.g. reference_time as ISO string) that would replace the # datetime values set above. for k, v in (edge.attributes or {}).items(): if k not in edge_data: edge_data[k] = v edges.append(edge_data) if driver.graph_operations_interface: await driver.graph_operations_interface.episodic_node_save_bulk(None, driver, tx, episodes) await driver.graph_operations_interface.node_save_bulk(None, driver, tx, nodes) await driver.graph_operations_interface.episodic_edge_save_bulk( None, driver, tx, [edge.model_dump() for edge in episodic_edges] ) await driver.graph_operations_interface.edge_save_bulk(None, driver, tx, edges) elif driver.provider == GraphProvider.KUZU: # FIXME: Kuzu's UNWIND does not currently support STRUCT[] type properly, so we insert the data one by one instead for now. episode_query = get_episode_node_save_bulk_query(driver.provider) for episode in episodes: await tx.run(episode_query, **episode) entity_node_query = get_entity_node_save_bulk_query(driver.provider, nodes) for node in nodes: await tx.run(entity_node_query, **node) entity_edge_query = get_entity_edge_save_bulk_query(driver.provider) for edge in edges: await tx.run(entity_edge_query, **edge) episodic_edge_query = get_episodic_edge_save_bulk_query(driver.provider) for edge in episodic_edges: await tx.run(episodic_edge_query, **edge.model_dump()) else: await tx.run(get_episode_node_save_bulk_query(driver.provider), episodes=episodes) await tx.run( get_entity_node_save_bulk_query(driver.provider, nodes), nodes=nodes, ) await tx.run( get_episodic_edge_save_bulk_query(driver.provider), episodic_edges=[edge.model_dump() for edge in episodic_edges], ) await tx.run( get_entity_edge_save_bulk_query(driver.provider), entity_edges=edges, ) async def extract_nodes_and_edges_bulk( clients: GraphitiClients, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]], edge_type_map: dict[tuple[str, str], list[str]], entity_types: dict[str, type[BaseModel]] | None = None, excluded_entity_types: list[str] | None = None, edge_types: dict[str, type[BaseModel]] | None = None, custom_extraction_instructions: str | None = None, use_combined_extraction: bool = False, ) -> tuple[list[list[EntityNode]], list[list[EntityEdge]]]: if use_combined_extraction: return await _extract_nodes_and_edges_bulk_combined( clients, episode_tuples, edge_type_map, entity_types=entity_types, excluded_entity_types=excluded_entity_types, edge_types=edge_types, custom_extraction_instructions=custom_extraction_instructions, ) return await _extract_nodes_and_edges_bulk_separate( clients, episode_tuples, edge_type_map, entity_types=entity_types, excluded_entity_types=excluded_entity_types, edge_types=edge_types, custom_extraction_instructions=custom_extraction_instructions, ) async def _extract_nodes_and_edges_bulk_combined( clients: GraphitiClients, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]], edge_type_map: dict[tuple[str, str], list[str]], entity_types: dict[str, type[BaseModel]] | None = None, excluded_entity_types: list[str] | None = None, edge_types: dict[str, type[BaseModel]] | None = None, custom_extraction_instructions: str | None = None, ) -> tuple[list[list[EntityNode]], list[list[EntityEdge]]]: """Combined extraction: single LLM call per episode for both nodes and edges.""" from graphiti_core.utils.maintenance.combined_extraction import ( extract_nodes_and_edges as extract_combined, ) results = await semaphore_gather( *[ extract_combined( clients, episode, previous_episodes, entity_types=entity_types, excluded_entity_types=excluded_entity_types, edge_type_map=edge_type_map, edge_types=edge_types, custom_extraction_instructions=custom_extraction_instructions, ) for episode, previous_episodes in episode_tuples ] ) nodes_bulk = [nodes for nodes, _, _ in results] edges_bulk = [edges for _, edges, _ in results] return nodes_bulk, edges_bulk async def _extract_nodes_and_edges_bulk_separate( clients: GraphitiClients, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]], edge_type_map: dict[tuple[str, str], list[str]], entity_types: dict[str, type[BaseModel]] | None = None, excluded_entity_types: list[str] | None = None, edge_types: dict[str, type[BaseModel]] | None = None, custom_extraction_instructions: str | None = None, ) -> tuple[list[list[EntityNode]], list[list[EntityEdge]]]: """Separate extraction: two sequential LLM calls per episode (legacy).""" extracted_results: list[tuple[list[EntityNode], dict[str, list[int]]]] = await semaphore_gather( *[ extract_nodes( clients, episode, previous_episodes, entity_types=entity_types, excluded_entity_types=excluded_entity_types, custom_extraction_instructions=custom_extraction_instructions, ) for episode, previous_episodes in episode_tuples ] ) extracted_nodes_bulk = [nodes for nodes, _ in extracted_results] extracted_edges_bulk: list[list[EntityEdge]] = await semaphore_gather( *[ extract_edges( clients, episode, extracted_nodes_bulk[i], previous_episodes, edge_type_map=edge_type_map, group_id=episode.group_id, edge_types=edge_types, custom_extraction_instructions=custom_extraction_instructions, ) for i, (episode, previous_episodes) in enumerate(episode_tuples) ] ) return extracted_nodes_bulk, extracted_edges_bulk async def dedupe_nodes_bulk( clients: GraphitiClients, extracted_nodes: list[list[EntityNode]], episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]], entity_types: dict[str, type[BaseModel]] | None = None, ) -> tuple[dict[str, list[EntityNode]], dict[str, str]]: """Resolve entity duplicates across an in-memory batch using a two-pass strategy. 1. Run :func:`resolve_extracted_nodes` for every episode in parallel so each batch item is reconciled against the live graph just like the non-batch flow. 2. Re-run the deterministic similarity heuristics across the union of resolved nodes to catch duplicates that only co-occur inside this batch, emitting a canonical UUID map that callers can apply to edges and persistence. """ first_pass_results = await semaphore_gather( *[ resolve_extracted_nodes( clients, nodes, episode_tuples[i][0], episode_tuples[i][1], entity_types, ) for i, nodes in enumerate(extracted_nodes) ] ) episode_resolutions: list[tuple[str, list[EntityNode]]] = [] per_episode_uuid_maps: list[dict[str, str]] = [] duplicate_pairs: list[tuple[str, str]] = [] for (resolved_nodes, uuid_map, duplicates), (episode, _) in zip( first_pass_results, episode_tuples, strict=True ): episode_resolutions.append((episode.uuid, resolved_nodes)) per_episode_uuid_maps.append(uuid_map) duplicate_pairs.extend((source.uuid, target.uuid) for source, target in duplicates) canonical_nodes: dict[str, EntityNode] = {} for _, resolved_nodes in episode_resolutions: for node in resolved_nodes: # NOTE: this loop is O(n^2) in the number of nodes inside the batch because we rebuild # the MinHash index for the accumulated canonical pool each time. The LRU-backed # shingle cache keeps the constant factors low for typical batch sizes (≤ CHUNK_SIZE), # but if batches grow significantly we should switch to an incremental index or chunked # processing. if not canonical_nodes: canonical_nodes[node.uuid] = node continue existing_candidates = list(canonical_nodes.values()) normalized = _normalize_string_exact(node.name) exact_match = next( ( candidate for candidate in existing_candidates if _normalize_string_exact(candidate.name) == normalized ), None, ) if exact_match is not None: if exact_match.uuid != node.uuid: duplicate_pairs.append((node.uuid, exact_match.uuid)) continue indexes = _build_candidate_indexes(existing_candidates) state = DedupResolutionState( resolved_nodes=[None], uuid_map={}, unresolved_indices=[], ) _resolve_with_similarity([node], indexes, state) resolved = state.resolved_nodes[0] if resolved is None: canonical_nodes[node.uuid] = node continue canonical_uuid = resolved.uuid canonical_nodes.setdefault(canonical_uuid, resolved) if canonical_uuid != node.uuid: duplicate_pairs.append((node.uuid, canonical_uuid)) union_pairs: list[tuple[str, str]] = [] for uuid_map in per_episode_uuid_maps: union_pairs.extend(uuid_map.items()) union_pairs.extend(duplicate_pairs) compressed_map: dict[str, str] = _build_directed_uuid_map(union_pairs) nodes_by_episode: dict[str, list[EntityNode]] = {} for episode_uuid, resolved_nodes in episode_resolutions: deduped_nodes: list[EntityNode] = [] seen: set[str] = set() for node in resolved_nodes: canonical_uuid = compressed_map.get(node.uuid, node.uuid) if canonical_uuid in seen: continue seen.add(canonical_uuid) canonical_node = canonical_nodes.get(canonical_uuid) if canonical_node is None: logger.error( 'Canonical node %s missing during batch dedupe; falling back to %s', canonical_uuid, node.uuid, ) canonical_node = node deduped_nodes.append(canonical_node) nodes_by_episode[episode_uuid] = deduped_nodes return nodes_by_episode, compressed_map async def dedupe_edges_bulk( clients: GraphitiClients, extracted_edges: list[list[EntityEdge]], episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]], _entities: list[EntityNode], edge_types: dict[str, type[BaseModel]], _edge_type_map: dict[tuple[str, str], list[str]], ) -> dict[str, list[EntityEdge]]: embedder = clients.embedder min_score = 0.6 # generate embeddings await semaphore_gather( *[create_entity_edge_embeddings(embedder, edges) for edges in extracted_edges] ) # Find similar results dedupe_tuples: list[tuple[EpisodicNode, EntityEdge, list[EntityEdge]]] = [] for i, edges_i in enumerate(extracted_edges): existing_edges: list[EntityEdge] = [] for edges_j in extracted_edges: existing_edges += edges_j for edge in edges_i: candidates: list[EntityEdge] = [] for existing_edge in existing_edges: # Skip self-comparison if edge.uuid == existing_edge.uuid: continue # Approximate BM25 by checking for word overlaps (this is faster than creating many in-memory indices) # This approach will cast a wider net than BM25, which is ideal for this use case if ( edge.source_node_uuid != existing_edge.source_node_uuid or edge.target_node_uuid != existing_edge.target_node_uuid ): continue edge_words = set(edge.fact.lower().split()) existing_edge_words = set(existing_edge.fact.lower().split()) has_overlap = not edge_words.isdisjoint(existing_edge_words) if has_overlap: candidates.append(existing_edge) continue # Check for semantic similarity even if there is no overlap similarity = np.dot( normalize_l2(edge.fact_embedding or []), normalize_l2(existing_edge.fact_embedding or []), ) if similarity >= min_score: candidates.append(existing_edge) dedupe_tuples.append((episode_tuples[i][0], edge, candidates)) bulk_edge_resolutions: list[ tuple[EntityEdge, EntityEdge, list[EntityEdge]] ] = await semaphore_gather( *[ resolve_extracted_edge( clients.llm_client, edge, candidates, candidates, episode, edge_types, ) for episode, edge, candidates in dedupe_tuples ] ) # For now we won't track edge invalidation duplicate_pairs: list[tuple[str, str]] = [] for i, (_, _, duplicates) in enumerate(bulk_edge_resolutions): episode, edge, candidates = dedupe_tuples[i] for duplicate in duplicates: duplicate_pairs.append((edge.uuid, duplicate.uuid)) # Now we compress the duplicate_map, so that 3 -> 2 and 2 -> becomes 3 -> 1 (sorted by uuid) compressed_map: dict[str, str] = compress_uuid_map(duplicate_pairs) edge_uuid_map: dict[str, EntityEdge] = { edge.uuid: edge for edges in extracted_edges for edge in edges } edges_by_episode: dict[str, list[EntityEdge]] = {} for i, edges in enumerate(extracted_edges): episode = episode_tuples[i][0] edges_by_episode[episode.uuid] = [ edge_uuid_map[compressed_map.get(edge.uuid, edge.uuid)] for edge in edges ] return edges_by_episode class UnionFind: def __init__(self, elements): # start each element in its own set self.parent = {e: e for e in elements} def find(self, x): # path‐compression if self.parent[x] != x: self.parent[x] = self.find(self.parent[x]) return self.parent[x] def union(self, a, b): ra, rb = self.find(a), self.find(b) if ra == rb: return # attach the lexicographically larger root under the smaller if ra < rb: self.parent[rb] = ra else: self.parent[ra] = rb def compress_uuid_map(duplicate_pairs: list[tuple[str, str]]) -> dict[str, str]: """ all_ids: iterable of all entity IDs (strings) duplicate_pairs: iterable of (id1, id2) pairs returns: dict mapping each id -> lexicographically smallest id in its duplicate set """ all_uuids = set() for pair in duplicate_pairs: all_uuids.add(pair[0]) all_uuids.add(pair[1]) uf = UnionFind(all_uuids) for a, b in duplicate_pairs: uf.union(a, b) # ensure full path‐compression before mapping return {uuid: uf.find(uuid) for uuid in all_uuids} E = typing.TypeVar('E', bound=Edge) def resolve_edge_pointers(edges: list[E], uuid_map: dict[str, str]): for edge in edges: source_uuid = edge.source_node_uuid target_uuid = edge.target_node_uuid edge.source_node_uuid = uuid_map.get(source_uuid, source_uuid) edge.target_node_uuid = uuid_map.get(target_uuid, target_uuid) return edges