graphiti-core 0.17.4__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/utils/maintenance/edge_operations.py

@@ -21,7 +21,7 @@ from time import time
 from pydantic import BaseModel
 from typing_extensions import LiteralString
 
-from graphiti_core.driver.driver import GraphDriver
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
 from graphiti_core.edges import (
     CommunityEdge,
     EntityEdge,
@@ -34,11 +34,16 @@ from graphiti_core.llm_client import LLMClient
 from graphiti_core.llm_client.config import ModelSize
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
 from graphiti_core.prompts import prompt_library
-from graphiti_core.prompts.dedupe_edges import EdgeDuplicate, UniqueFacts
+from graphiti_core.prompts.dedupe_edges import EdgeDuplicate
 from graphiti_core.prompts.extract_edges import ExtractedEdges, MissingFacts
+from graphiti_core.search.search import search
+from graphiti_core.search.search_config import SearchResults
+from graphiti_core.search.search_config_recipes import EDGE_HYBRID_SEARCH_RRF
 from graphiti_core.search.search_filters import SearchFilters
-from graphiti_core.search.search_utils import get_edge_invalidation_candidates, get_relevant_edges
 from graphiti_core.utils.datetime_utils import ensure_utc, utc_now
+from graphiti_core.utils.maintenance.dedup_helpers import _normalize_string_exact
+
+DEFAULT_EDGE_NAME = 'RELATES_TO'
 
 logger = logging.getLogger(__name__)
 
@@ -63,32 +68,6 @@ def build_episodic_edges(
     return episodic_edges
 
 
-def build_duplicate_of_edges(
-    episode: EpisodicNode,
-    created_at: datetime,
-    duplicate_nodes: list[tuple[EntityNode, EntityNode]],
-) -> list[EntityEdge]:
-    is_duplicate_of_edges: list[EntityEdge] = []
-    for source_node, target_node in duplicate_nodes:
-        if source_node.uuid == target_node.uuid:
-            continue
-
-        is_duplicate_of_edges.append(
-            EntityEdge(
-                source_node_uuid=source_node.uuid,
-                target_node_uuid=target_node.uuid,
-                name='IS_DUPLICATE_OF',
-                group_id=episode.group_id,
-                fact=f'{source_node.name} is a duplicate of {target_node.name}',
-                episodes=[episode.uuid],
-                created_at=created_at,
-                valid_at=created_at,
-            )
-        )
-
-    return is_duplicate_of_edges
-
-
 def build_community_edges(
     entity_nodes: list[EntityNode],
     community_node: CommunityNode,
@@ -114,7 +93,7 @@ async def extract_edges(
     previous_episodes: list[EpisodicNode],
     edge_type_map: dict[tuple[str, str], list[str]],
     group_id: str = '',
-    edge_types: dict[str, BaseModel] | None = None,
+    edge_types: dict[str, type[BaseModel]] | None = None,
 ) -> list[EntityEdge]:
     start = time()
 
@@ -160,10 +139,12 @@ async def extract_edges(
             prompt_library.extract_edges.edge(context),
             response_model=ExtractedEdges,
             max_tokens=extract_edges_max_tokens,
+            group_id=group_id,
+            prompt_name='extract_edges.edge',
         )
-        edges_data = llm_response.get('edges', [])
+        edges_data = ExtractedEdges(**llm_response).edges
 
-        context['extracted_facts'] = [edge_data.get('fact', '') for edge_data in edges_data]
+        context['extracted_facts'] = [edge_data.fact for edge_data in edges_data]
 
         reflexion_iterations += 1
         if reflexion_iterations < MAX_REFLEXION_ITERATIONS:
@@ -171,6 +152,8 @@ async def extract_edges(
                 prompt_library.extract_edges.reflexion(context),
                 response_model=MissingFacts,
                 max_tokens=extract_edges_max_tokens,
+                group_id=group_id,
+                prompt_name='extract_edges.reflexion',
             )
 
             missing_facts = reflexion_response.get('missing_facts', [])
@@ -193,20 +176,31 @@ async def extract_edges(
     edges = []
     for edge_data in edges_data:
         # Validate Edge Date information
-        valid_at = edge_data.get('valid_at', None)
-        invalid_at = edge_data.get('invalid_at', None)
+        valid_at = edge_data.valid_at
+        invalid_at = edge_data.invalid_at
         valid_at_datetime = None
         invalid_at_datetime = None
 
-        source_node_idx = edge_data.get('source_entity_id', -1)
-        target_node_idx = edge_data.get('target_entity_id', -1)
-        if not (-1 < source_node_idx < len(nodes) and -1 < target_node_idx < len(nodes)):
+        # Filter out empty edges
+        if not edge_data.fact.strip():
+            continue
+
+        source_node_idx = edge_data.source_entity_id
+        target_node_idx = edge_data.target_entity_id
+
+        if len(nodes) == 0:
+            logger.warning('No entities provided for edge extraction')
+            continue
+
+        if not (0 <= source_node_idx < len(nodes) and 0 <= target_node_idx < len(nodes)):
             logger.warning(
-                f'WARNING: source or target node not filled {edge_data.get("edge_name")}. source_node_uuid: {source_node_idx} and target_node_uuid: {target_node_idx} '
+                f'Invalid entity IDs in edge extraction for {edge_data.relation_type}. '
+                f'source_entity_id: {source_node_idx}, target_entity_id: {target_node_idx}, '
+                f'but only {len(nodes)} entities available (valid range: 0-{len(nodes) - 1})'
             )
             continue
         source_node_uuid = nodes[source_node_idx].uuid
-        target_node_uuid = nodes[edge_data.get('target_entity_id')].uuid
+        target_node_uuid = nodes[target_node_idx].uuid
 
         if valid_at:
             try:
@@ -226,9 +220,9 @@ async def extract_edges(
         edge = EntityEdge(
             source_node_uuid=source_node_uuid,
             target_node_uuid=target_node_uuid,
-            name=edge_data.get('relation_type', ''),
+            name=edge_data.relation_type,
             group_id=group_id,
-            fact=edge_data.get('fact', ''),
+            fact=edge_data.fact,
             episodes=[episode.uuid],
             created_at=utc_now(),
             valid_at=valid_at_datetime,
@@ -249,20 +243,68 @@ async def resolve_extracted_edges(
     extracted_edges: list[EntityEdge],
     episode: EpisodicNode,
     entities: list[EntityNode],
-    edge_types: dict[str, BaseModel],
+    edge_types: dict[str, type[BaseModel]],
     edge_type_map: dict[tuple[str, str], list[str]],
 ) -> tuple[list[EntityEdge], list[EntityEdge]]:
+    # Fast path: deduplicate exact matches within the extracted edges before parallel processing
+    seen: dict[tuple[str, str, str], EntityEdge] = {}
+    deduplicated_edges: list[EntityEdge] = []
+
+    for edge in extracted_edges:
+        key = (
+            edge.source_node_uuid,
+            edge.target_node_uuid,
+            _normalize_string_exact(edge.fact),
+        )
+        if key not in seen:
+            seen[key] = edge
+            deduplicated_edges.append(edge)
+
+    extracted_edges = deduplicated_edges
+
     driver = clients.driver
     llm_client = clients.llm_client
     embedder = clients.embedder
     await create_entity_edge_embeddings(embedder, extracted_edges)
 
-    search_results = await semaphore_gather(
-        get_relevant_edges(driver, extracted_edges, SearchFilters()),
-        get_edge_invalidation_candidates(driver, extracted_edges, SearchFilters(), 0.2),
+    valid_edges_list: list[list[EntityEdge]] = await semaphore_gather(
+        *[
+            EntityEdge.get_between_nodes(driver, edge.source_node_uuid, edge.target_node_uuid)
+            for edge in extracted_edges
+        ]
+    )
+
+    related_edges_results: list[SearchResults] = await semaphore_gather(
+        *[
+            search(
+                clients,
+                extracted_edge.fact,
+                group_ids=[extracted_edge.group_id],
+                config=EDGE_HYBRID_SEARCH_RRF,
+                search_filter=SearchFilters(edge_uuids=[edge.uuid for edge in valid_edges]),
+            )
+            for extracted_edge, valid_edges in zip(extracted_edges, valid_edges_list, strict=True)
+        ]
     )
 
-    related_edges_lists, edge_invalidation_candidates = search_results
+    related_edges_lists: list[list[EntityEdge]] = [result.edges for result in related_edges_results]
+
+    edge_invalidation_candidate_results: list[SearchResults] = await semaphore_gather(
+        *[
+            search(
+                clients,
+                extracted_edge.fact,
+                group_ids=[extracted_edge.group_id],
+                config=EDGE_HYBRID_SEARCH_RRF,
+                search_filter=SearchFilters(),
+            )
+            for extracted_edge in extracted_edges
+        ]
+    )
+
+    edge_invalidation_candidates: list[list[EntityEdge]] = [
+        result.edges for result in edge_invalidation_candidate_results
+    ]
 
     logger.debug(
         f'Related edges lists: {[(e.name, e.uuid) for edges_lst in related_edges_lists for e in edges_lst]}'
@@ -271,11 +313,21 @@ async def resolve_extracted_edges(
     # Build entity hash table
     uuid_entity_map: dict[str, EntityNode] = {entity.uuid: entity for entity in entities}
 
-    # Determine which edge types are relevant for each edge
-    edge_types_lst: list[dict[str, BaseModel]] = []
+    # Determine which edge types are relevant for each edge.
+    # `edge_types_lst` stores the subset of custom edge definitions whose
+    # node signature matches each extracted edge. Anything outside this subset
+    # should only stay on the edge if it is a non-custom (LLM generated) label.
+    edge_types_lst: list[dict[str, type[BaseModel]]] = []
+    custom_type_names = set(edge_types or {})
     for extracted_edge in extracted_edges:
-        source_node_labels = uuid_entity_map[extracted_edge.source_node_uuid].labels + ['Entity']
-        target_node_labels = uuid_entity_map[extracted_edge.target_node_uuid].labels + ['Entity']
+        source_node = uuid_entity_map.get(extracted_edge.source_node_uuid)
+        target_node = uuid_entity_map.get(extracted_edge.target_node_uuid)
+        source_node_labels = (
+            source_node.labels + ['Entity'] if source_node is not None else ['Entity']
+        )
+        target_node_labels = (
+            target_node.labels + ['Entity'] if target_node is not None else ['Entity']
+        )
         label_tuples = [
             (source_label, target_label)
             for source_label in source_node_labels
@@ -294,6 +346,20 @@ async def resolve_extracted_edges(
 
         edge_types_lst.append(extracted_edge_types)
 
+    for extracted_edge, extracted_edge_types in zip(extracted_edges, edge_types_lst, strict=True):
+        allowed_type_names = set(extracted_edge_types)
+        is_custom_name = extracted_edge.name in custom_type_names
+        if not allowed_type_names:
+            # No custom types are valid for this node pairing. Keep LLM generated
+            # labels, but flip disallowed custom names back to the default.
+            if is_custom_name and extracted_edge.name != DEFAULT_EDGE_NAME:
+                extracted_edge.name = DEFAULT_EDGE_NAME
+            continue
+        if is_custom_name and extracted_edge.name not in allowed_type_names:
+            # Custom name exists but it is not permitted for this source/target
+            # signature, so fall back to the default edge label.
+            extracted_edge.name = DEFAULT_EDGE_NAME
+
     # resolve edges with related edges in the graph and find invalidation candidates
     results: list[tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]] = list(
         await semaphore_gather(
@@ -305,6 +371,7 @@ async def resolve_extracted_edges(
                     existing_edges,
                     episode,
                     extracted_edge_types,
+                    custom_type_names,
                 )
                 for extracted_edge, related_edges, existing_edges, extracted_edge_types in zip(
                     extracted_edges,
@@ -346,21 +413,26 @@ def resolve_edge_contradictions(
     invalidated_edges: list[EntityEdge] = []
     for edge in invalidation_candidates:
         # (Edge invalid before new edge becomes valid) or (new edge invalid before edge becomes valid)
+        edge_invalid_at_utc = ensure_utc(edge.invalid_at)
+        resolved_edge_valid_at_utc = ensure_utc(resolved_edge.valid_at)
+        edge_valid_at_utc = ensure_utc(edge.valid_at)
+        resolved_edge_invalid_at_utc = ensure_utc(resolved_edge.invalid_at)
+
         if (
-            edge.invalid_at is not None
-            and resolved_edge.valid_at is not None
-            and edge.invalid_at <= resolved_edge.valid_at
+            edge_invalid_at_utc is not None
+            and resolved_edge_valid_at_utc is not None
+            and edge_invalid_at_utc <= resolved_edge_valid_at_utc
         ) or (
-            edge.valid_at is not None
-            and resolved_edge.invalid_at is not None
-            and resolved_edge.invalid_at <= edge.valid_at
+            edge_valid_at_utc is not None
+            and resolved_edge_invalid_at_utc is not None
+            and resolved_edge_invalid_at_utc <= edge_valid_at_utc
         ):
             continue
         # New edge invalidates edge
         elif (
-            edge.valid_at is not None
-            and resolved_edge.valid_at is not None
-            and edge.valid_at < resolved_edge.valid_at
+            edge_valid_at_utc is not None
+            and resolved_edge_valid_at_utc is not None
+            and edge_valid_at_utc < resolved_edge_valid_at_utc
         ):
             edge.invalid_at = resolved_edge.valid_at
             edge.expired_at = edge.expired_at if edge.expired_at is not None else utc_now()
@@ -375,32 +447,69 @@ async def resolve_extracted_edge(
     related_edges: list[EntityEdge],
     existing_edges: list[EntityEdge],
     episode: EpisodicNode,
-    edge_types: dict[str, BaseModel] | None = None,
+    edge_type_candidates: dict[str, type[BaseModel]] | None = None,
+    custom_edge_type_names: set[str] | None = None,
 ) -> tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]:
+    """Resolve an extracted edge against existing graph context.
+
+    Parameters
+    ----------
+    llm_client : LLMClient
+        Client used to invoke the LLM for deduplication and attribute extraction.
+    extracted_edge : EntityEdge
+        Newly extracted edge whose canonical representation is being resolved.
+    related_edges : list[EntityEdge]
+        Candidate edges with identical endpoints used for duplicate detection.
+    existing_edges : list[EntityEdge]
+        Broader set of edges evaluated for contradiction / invalidation.
+    episode : EpisodicNode
+        Episode providing content context when extracting edge attributes.
+    edge_type_candidates : dict[str, type[BaseModel]] | None
+        Custom edge types permitted for the current source/target signature.
+    custom_edge_type_names : set[str] | None
+        Full catalog of registered custom edge names. Used to distinguish
+        between disallowed custom types (which fall back to the default label)
+        and ad-hoc labels emitted by the LLM.
+
+    Returns
+    -------
+    tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]
+        The resolved edge, any duplicates, and edges to invalidate.
+    """
     if len(related_edges) == 0 and len(existing_edges) == 0:
         return extracted_edge, [], []
 
+    # Fast path: if the fact text and endpoints already exist verbatim, reuse the matching edge.
+    normalized_fact = _normalize_string_exact(extracted_edge.fact)
+    for edge in related_edges:
+        if (
+            edge.source_node_uuid == extracted_edge.source_node_uuid
+            and edge.target_node_uuid == extracted_edge.target_node_uuid
+            and _normalize_string_exact(edge.fact) == normalized_fact
+        ):
+            resolved = edge
+            if episode is not None and episode.uuid not in resolved.episodes:
+                resolved.episodes.append(episode.uuid)
+            return resolved, [], []
+
     start = time()
 
     # Prepare context for LLM
-    related_edges_context = [
-        {'id': edge.uuid, 'fact': edge.fact} for i, edge in enumerate(related_edges)
-    ]
+    related_edges_context = [{'idx': i, 'fact': edge.fact} for i, edge in enumerate(related_edges)]
 
     invalidation_edge_candidates_context = [
-        {'id': i, 'fact': existing_edge.fact} for i, existing_edge in enumerate(existing_edges)
+        {'idx': i, 'fact': existing_edge.fact} for i, existing_edge in enumerate(existing_edges)
     ]
 
     edge_types_context = (
         [
             {
-                'fact_type_id': i,
                 'fact_type_name': type_name,
                 'fact_type_description': type_model.__doc__,
             }
-            for i, (type_name, type_model) in enumerate(edge_types.items())
+            for type_name, type_model in edge_type_candidates.items()
         ]
-        if edge_types is not None
+        if edge_type_candidates is not None
         else []
     )
 
@@ -411,15 +520,34 @@ async def resolve_extracted_edge(
         'edge_types': edge_types_context,
     }
 
+    if related_edges or existing_edges:
+        logger.debug(
+            'Resolving edge: sent %d EXISTING FACTS%s and %d INVALIDATION CANDIDATES%s',
+            len(related_edges),
+            f' (idx 0-{len(related_edges) - 1})' if related_edges else '',
+            len(existing_edges),
+            f' (idx 0-{len(existing_edges) - 1})' if existing_edges else '',
+        )
+
     llm_response = await llm_client.generate_response(
         prompt_library.dedupe_edges.resolve_edge(context),
         response_model=EdgeDuplicate,
         model_size=ModelSize.small,
+        prompt_name='dedupe_edges.resolve_edge',
     )
+    response_object = EdgeDuplicate(**llm_response)
+    duplicate_facts = response_object.duplicate_facts
+
+    # Validate duplicate_facts are in valid range for EXISTING FACTS
+    invalid_duplicates = [i for i in duplicate_facts if i < 0 or i >= len(related_edges)]
+    if invalid_duplicates:
+        logger.warning(
+            'LLM returned invalid duplicate_facts idx values %s (valid range: 0-%d for EXISTING FACTS)',
+            invalid_duplicates,
+            len(related_edges) - 1,
+        )
 
-    duplicate_fact_ids: list[int] = list(
-        filter(lambda i: 0 <= i < len(related_edges), llm_response.get('duplicate_facts', []))
-    )
+    duplicate_fact_ids: list[int] = [i for i in duplicate_facts if 0 <= i < len(related_edges)]
 
     resolved_edge = extracted_edge
     for duplicate_fact_id in duplicate_fact_ids:
@@ -429,12 +557,32 @@ async def resolve_extracted_edge(
     if duplicate_fact_ids and episode is not None:
         resolved_edge.episodes.append(episode.uuid)
 
-    contradicted_facts: list[int] = llm_response.get('contradicted_facts', [])
+    contradicted_facts: list[int] = response_object.contradicted_facts
 
-    invalidation_candidates: list[EntityEdge] = [existing_edges[i] for i in contradicted_facts]
+    # Validate contradicted_facts are in valid range for INVALIDATION CANDIDATES
+    invalid_contradictions = [i for i in contradicted_facts if i < 0 or i >= len(existing_edges)]
+    if invalid_contradictions:
+        logger.warning(
+            'LLM returned invalid contradicted_facts idx values %s (valid range: 0-%d for INVALIDATION CANDIDATES)',
+            invalid_contradictions,
+            len(existing_edges) - 1,
+        )
 
-    fact_type: str = str(llm_response.get('fact_type'))
-    if fact_type.upper() != 'DEFAULT' and edge_types is not None:
+    invalidation_candidates: list[EntityEdge] = [
+        existing_edges[i] for i in contradicted_facts if 0 <= i < len(existing_edges)
+    ]
+
+    fact_type: str = response_object.fact_type
+    candidate_type_names = set(edge_type_candidates or {})
+    custom_type_names = custom_edge_type_names or set()
+
+    is_default_type = fact_type.upper() == 'DEFAULT'
+    is_custom_type = fact_type in custom_type_names
+    is_allowed_custom_type = fact_type in candidate_type_names
+
+    if is_allowed_custom_type:
+        # The LLM selected a custom type that is allowed for the node pair.
+        # Adopt the custom type and, if needed, extract its structured attributes.
         resolved_edge.name = fact_type
 
         edge_attributes_context = {
@@ -443,15 +591,26 @@ async def resolve_extracted_edge(
             'fact': resolved_edge.fact,
         }
 
-        edge_model = edge_types.get(fact_type)
+        edge_model = edge_type_candidates.get(fact_type) if edge_type_candidates else None
         if edge_model is not None and len(edge_model.model_fields) != 0:
             edge_attributes_response = await llm_client.generate_response(
                 prompt_library.extract_edges.extract_attributes(edge_attributes_context),
                 response_model=edge_model,  # type: ignore
                 model_size=ModelSize.small,
+                prompt_name='extract_edges.extract_attributes',
             )
 
             resolved_edge.attributes = edge_attributes_response
+    elif not is_default_type and is_custom_type:
+        # The LLM picked a custom type that is not allowed for this signature.
+        # Reset to the default label and drop any structured attributes.
+        resolved_edge.name = DEFAULT_EDGE_NAME
+        resolved_edge.attributes = {}
+    elif not is_default_type:
+        # Non-custom labels are allowed to pass through so long as the LLM does
+        # not return the sentinel DEFAULT value.
+        resolved_edge.name = fact_type
+        resolved_edge.attributes = {}
 
     end = time()
     logger.debug(
@@ -465,14 +624,14 @@ async def resolve_extracted_edge(
 
     # Determine if the new_edge needs to be expired
     if resolved_edge.expired_at is None:
-        invalidation_candidates.sort(key=lambda c: (c.valid_at is None, c.valid_at))
+        invalidation_candidates.sort(key=lambda c: (c.valid_at is None, ensure_utc(c.valid_at)))
         for candidate in invalidation_candidates:
+            candidate_valid_at_utc = ensure_utc(candidate.valid_at)
+            resolved_edge_valid_at_utc = ensure_utc(resolved_edge.valid_at)
             if (
-                candidate.valid_at
-                and resolved_edge.valid_at
-                and candidate.valid_at.tzinfo
-                and resolved_edge.valid_at.tzinfo
-                and candidate.valid_at > resolved_edge.valid_at
+                candidate_valid_at_utc is not None
+                and resolved_edge_valid_at_utc is not None
+                and candidate_valid_at_utc > resolved_edge_valid_at_utc
             ):
                 # Expire new edge since we have information about more recent events
                 resolved_edge.invalid_at = candidate.valid_at
@@ -488,59 +647,60 @@ async def resolve_extracted_edge(
     return resolved_edge, invalidated_edges, duplicate_edges
 
 
-async def dedupe_edge_list(
-    llm_client: LLMClient,
-    edges: list[EntityEdge],
-) -> list[EntityEdge]:
-    start = time()
-
-    # Create edge map
-    edge_map = {}
-    for edge in edges:
-        edge_map[edge.uuid] = edge
-
-    # Prepare context for LLM
-    context = {'edges': [{'uuid': edge.uuid, 'fact': edge.fact} for edge in edges]}
-
-    llm_response = await llm_client.generate_response(
-        prompt_library.dedupe_edges.edge_list(context), response_model=UniqueFacts
-    )
-    unique_edges_data = llm_response.get('unique_facts', [])
-
-    end = time()
-    logger.debug(f'Extracted edge duplicates: {unique_edges_data} in {(end - start) * 1000} ms ')
-
-    # Get full edge data
-    unique_edges = []
-    for edge_data in unique_edges_data:
-        uuid = edge_data['uuid']
-        edge = edge_map[uuid]
-        edge.fact = edge_data['fact']
-        unique_edges.append(edge)
-
-    return unique_edges
-
-
 async def filter_existing_duplicate_of_edges(
     driver: GraphDriver, duplicates_node_tuples: list[tuple[EntityNode, EntityNode]]
 ) -> list[tuple[EntityNode, EntityNode]]:
-    query: LiteralString = """
-        UNWIND $duplicate_node_uuids AS duplicate_tuple
-        MATCH (n:Entity {uuid: duplicate_tuple[0]})-[r:RELATES_TO {name: 'IS_DUPLICATE_OF'}]->(m:Entity {uuid: duplicate_tuple[1]})
-        RETURN DISTINCT
-            n.uuid AS source_uuid,
-            m.uuid AS target_uuid
-    """
+    if not duplicates_node_tuples:
+        return []
 
     duplicate_nodes_map = {
         (source.uuid, target.uuid): (source, target) for source, target in duplicates_node_tuples
     }
 
-    records, _, _ = await driver.execute_query(
-        query,
-        duplicate_node_uuids=list(duplicate_nodes_map.keys()),
-        routing_='r',
-    )
+    if driver.provider == GraphProvider.NEPTUNE:
+        query: LiteralString = """
+            UNWIND $duplicate_node_uuids AS duplicate_tuple
+            MATCH (n:Entity {uuid: duplicate_tuple.source})-[r:RELATES_TO {name: 'IS_DUPLICATE_OF'}]->(m:Entity {uuid: duplicate_tuple.target})
+            RETURN DISTINCT
+                n.uuid AS source_uuid,
+                m.uuid AS target_uuid
+        """
+
+        duplicate_nodes = [
+            {'source': source.uuid, 'target': target.uuid}
+            for source, target in duplicates_node_tuples
+        ]
+
+        records, _, _ = await driver.execute_query(
+            query,
+            duplicate_node_uuids=duplicate_nodes,
+            routing_='r',
+        )
+    else:
+        if driver.provider == GraphProvider.KUZU:
+            query = """
+                UNWIND $duplicate_node_uuids AS duplicate
+                MATCH (n:Entity {uuid: duplicate.src})-[:RELATES_TO]->(e:RelatesToNode_ {name: 'IS_DUPLICATE_OF'})-[:RELATES_TO]->(m:Entity {uuid: duplicate.dst})
+                RETURN DISTINCT
+                    n.uuid AS source_uuid,
+                    m.uuid AS target_uuid
+            """
+            duplicate_node_uuids = [{'src': src, 'dst': dst} for src, dst in duplicate_nodes_map]
+        else:
+            query: LiteralString = """
+                UNWIND $duplicate_node_uuids AS duplicate_tuple
+                MATCH (n:Entity {uuid: duplicate_tuple[0]})-[r:RELATES_TO {name: 'IS_DUPLICATE_OF'}]->(m:Entity {uuid: duplicate_tuple[1]})
+                RETURN DISTINCT
+                    n.uuid AS source_uuid,
+                    m.uuid AS target_uuid
+            """
+            duplicate_node_uuids = list(duplicate_nodes_map.keys())
+
+        records, _, _ = await driver.execute_query(
+            query,
+            duplicate_node_uuids=duplicate_node_uuids,
+            routing_='r',
+        )
 
     # Remove duplicates that already have the IS_DUPLICATE_OF edge
     for record in records: