graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/utils/maintenance/edge_operations.py

@@ -19,7 +19,9 @@ from datetime import datetime
 from time import time
 
 from pydantic import BaseModel
+from typing_extensions import LiteralString
 
+from graphiti_core.driver.driver import GraphDriver, GraphProvider
 from graphiti_core.edges import (
     CommunityEdge,
     EntityEdge,
@@ -32,26 +34,31 @@ 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__)
 
 
 def build_episodic_edges(
     entity_nodes: list[EntityNode],
-    episode: EpisodicNode,
+    episode_uuid: str,
     created_at: datetime,
 ) -> list[EpisodicEdge]:
     episodic_edges: list[EpisodicEdge] = [
         EpisodicEdge(
-            source_node_uuid=episode.uuid,
+            source_node_uuid=episode_uuid,
             target_node_uuid=node.uuid,
             created_at=created_at,
-            group_id=episode.group_id,
+            group_id=node.group_id,
         )
         for node in entity_nodes
     ]
@@ -84,20 +91,26 @@ async def extract_edges(
     episode: EpisodicNode,
     nodes: list[EntityNode],
     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()
 
     extract_edges_max_tokens = 16384
     llm_client = clients.llm_client
 
-    node_uuids_by_name_map = {node.name: node.uuid for node in nodes}
+    edge_type_signature_map: dict[str, tuple[str, str]] = {
+        edge_type: signature
+        for signature, edge_types in edge_type_map.items()
+        for edge_type in edge_types
+    }
 
     edge_types_context = (
         [
             {
                 'fact_type_name': type_name,
+                'fact_type_signature': edge_type_signature_map.get(type_name, ('Entity', 'Entity')),
                 'fact_type_description': type_model.__doc__,
             }
             for type_name, type_model in edge_types.items()
@@ -109,7 +122,10 @@ async def extract_edges(
     # Prepare context for LLM
     context = {
         'episode_content': episode.content,
-        'nodes': [node.name for node in nodes],
+        'nodes': [
+            {'id': idx, 'name': node.name, 'entity_types': node.labels}
+            for idx, node in enumerate(nodes)
+        ],
         'previous_episodes': [ep.content for ep in previous_episodes],
         'reference_time': episode.valid_at,
         'edge_types': edge_types_context,
@@ -123,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:
@@ -134,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', [])
@@ -156,11 +176,32 @@ 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
 
+        # 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'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[target_node_idx].uuid
+
         if valid_at:
             try:
                 valid_at_datetime = ensure_utc(
@@ -177,15 +218,11 @@ async def extract_edges(
             except ValueError as e:
                 logger.warning(f'WARNING: Error parsing invalid_at date: {e}. Input: {invalid_at}')
         edge = EntityEdge(
-            source_node_uuid=node_uuids_by_name_map.get(
-                edge_data.get('source_entity_name', ''), ''
-            ),
-            target_node_uuid=node_uuids_by_name_map.get(
-                edge_data.get('target_entity_name', ''), ''
-            ),
-            name=edge_data.get('relation_type', ''),
+            source_node_uuid=source_node_uuid,
+            target_node_uuid=target_node_uuid,
+            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,
@@ -201,70 +238,73 @@ async def extract_edges(
     return edges
 
 
-async def dedupe_extracted_edges(
-    llm_client: LLMClient,
-    extracted_edges: list[EntityEdge],
-    existing_edges: list[EntityEdge],
-) -> list[EntityEdge]:
-    # Create edge map
-    edge_map: dict[str, EntityEdge] = {}
-    for edge in existing_edges:
-        edge_map[edge.uuid] = edge
-
-    # Prepare context for LLM
-    context = {
-        'extracted_edges': [
-            {'uuid': edge.uuid, 'name': edge.name, 'fact': edge.fact} for edge in extracted_edges
-        ],
-        'existing_edges': [
-            {'uuid': edge.uuid, 'name': edge.name, 'fact': edge.fact} for edge in existing_edges
-        ],
-    }
-
-    llm_response = await llm_client.generate_response(prompt_library.dedupe_edges.edge(context))
-    duplicate_data = llm_response.get('duplicates', [])
-    logger.debug(f'Extracted unique edges: {duplicate_data}')
-
-    duplicate_uuid_map: dict[str, str] = {}
-    for duplicate in duplicate_data:
-        uuid_value = duplicate['duplicate_of']
-        duplicate_uuid_map[duplicate['uuid']] = uuid_value
-
-    # Get full edge data
-    edges: list[EntityEdge] = []
-    for edge in extracted_edges:
-        if edge.uuid in duplicate_uuid_map:
-            existing_uuid = duplicate_uuid_map[edge.uuid]
-            existing_edge = edge_map[existing_uuid]
-            # Add current episode to the episodes list
-            existing_edge.episodes += edge.episodes
-            edges.append(existing_edge)
-        else:
-            edges.append(edge)
-
-    return edges
-
-
 async def resolve_extracted_edges(
     clients: GraphitiClients,
     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: tuple[list[list[EntityEdge]], list[list[EntityEdge]]] = 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]}'
@@ -273,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
-        target_node_labels = uuid_entity_map[extracted_edge.target_node_uuid].labels
+        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
@@ -296,8 +346,22 @@ 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(
+    results: list[tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]] = list(
         await semaphore_gather(
             *[
                 resolve_extracted_edge(
@@ -307,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,
@@ -348,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()
@@ -377,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,
-) -> tuple[EntityEdge, list[EntityEdge]]:
+    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, []
+        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 []
     )
 
@@ -413,46 +520,97 @@ 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_id: int = llm_response.get('duplicate_fact_id', -1)
+    duplicate_fact_ids: list[int] = [i for i in duplicate_facts if 0 <= i < len(related_edges)]
 
-    resolved_edge = (
-        related_edges[duplicate_fact_id]
-        if 0 <= duplicate_fact_id < len(related_edges)
-        else extracted_edge
-    )
+    resolved_edge = extracted_edge
+    for duplicate_fact_id in duplicate_fact_ids:
+        resolved_edge = related_edges[duplicate_fact_id]
+        break
 
-    if duplicate_fact_id >= 0 and episode is not None:
+    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 = {
-            'message': episode.content,
+            'episode_content': episode.content,
             'reference_time': episode.valid_at,
             'fact': resolved_edge.fact,
         }
 
-        edge_model = edge_types.get(fact_type)
-
-        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,
-        )
+        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
+            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(
@@ -466,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
@@ -481,89 +639,73 @@ async def resolve_extracted_edge(
                 break
 
     # Determine which contradictory edges need to be expired
-    invalidated_edges = resolve_edge_contradictions(resolved_edge, invalidation_candidates)
-
-    return resolved_edge, invalidated_edges
-
-
-async def dedupe_extracted_edge(
-    llm_client: LLMClient,
-    extracted_edge: EntityEdge,
-    related_edges: list[EntityEdge],
-    episode: EpisodicNode | None = None,
-) -> EntityEdge:
-    if len(related_edges) == 0:
-        return extracted_edge
-
-    start = time()
-
-    # Prepare context for LLM
-    related_edges_context = [
-        {'id': edge.uuid, 'fact': edge.fact} for i, edge in enumerate(related_edges)
-    ]
-
-    extracted_edge_context = {
-        'fact': extracted_edge.fact,
-    }
-
-    context = {
-        'related_edges': related_edges_context,
-        'extracted_edges': extracted_edge_context,
-    }
-
-    llm_response = await llm_client.generate_response(
-        prompt_library.dedupe_edges.edge(context),
-        response_model=EdgeDuplicate,
-        model_size=ModelSize.small,
-    )
-
-    duplicate_fact_id: int = llm_response.get('duplicate_fact_id', -1)
-
-    edge = (
-        related_edges[duplicate_fact_id]
-        if 0 <= duplicate_fact_id < len(related_edges)
-        else extracted_edge
-    )
-
-    if duplicate_fact_id >= 0 and episode is not None:
-        edge.episodes.append(episode.uuid)
-
-    end = time()
-    logger.debug(
-        f'Resolved Edge: {extracted_edge.name} is {edge.name}, in {(end - start) * 1000} ms'
+    invalidated_edges: list[EntityEdge] = resolve_edge_contradictions(
+        resolved_edge, invalidation_candidates
     )
+    duplicate_edges: list[EntityEdge] = [related_edges[idx] for idx in duplicate_fact_ids]
 
-    return 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
+async def filter_existing_duplicate_of_edges(
+    driver: GraphDriver, duplicates_node_tuples: list[tuple[EntityNode, EntityNode]]
+) -> list[tuple[EntityNode, EntityNode]]:
+    if not duplicates_node_tuples:
+        return []
 
-    # Prepare context for LLM
-    context = {'edges': [{'uuid': edge.uuid, 'fact': edge.fact} for edge in edges]}
+    duplicate_nodes_map = {
+        (source.uuid, target.uuid): (source, target) for source, target in duplicates_node_tuples
+    }
 
-    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', [])
+    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
+        ]
 
-    end = time()
-    logger.debug(f'Extracted edge duplicates: {unique_edges_data} in {(end - start) * 1000} ms ')
+        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',
+        )
 
-    # 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)
+    # Remove duplicates that already have the IS_DUPLICATE_OF edge
+    for record in records:
+        duplicate_tuple = (record.get('source_uuid'), record.get('target_uuid'))
+        if duplicate_nodes_map.get(duplicate_tuple):
+            duplicate_nodes_map.pop(duplicate_tuple)
 
-    return unique_edges
+    return list(duplicate_nodes_map.values())