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

graphiti_core/utils/maintenance/node_operations.py

@@ -15,22 +15,26 @@ limitations under the License.
 """
 
 import logging
-from contextlib import suppress
+from collections.abc import Awaitable, Callable
 from time import time
 from typing import Any
-from uuid import uuid4
 
-import pydantic
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.llm_client.config import ModelSize
-from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
+from graphiti_core.nodes import (
+    EntityNode,
+    EpisodeType,
+    EpisodicNode,
+    create_entity_node_embeddings,
+)
 from graphiti_core.prompts import prompt_library
-from graphiti_core.prompts.dedupe_nodes import NodeResolutions
+from graphiti_core.prompts.dedupe_nodes import NodeDuplicate, NodeResolutions
 from graphiti_core.prompts.extract_nodes import (
+    EntitySummary,
     ExtractedEntities,
     ExtractedEntity,
     MissedEntities,
@@ -40,16 +44,28 @@ from graphiti_core.search.search_config import SearchResults
 from graphiti_core.search.search_config_recipes import NODE_HYBRID_SEARCH_RRF
 from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.utils.datetime_utils import utc_now
-from graphiti_core.utils.maintenance.edge_operations import filter_existing_duplicate_of_edges
+from graphiti_core.utils.maintenance.dedup_helpers import (
+    DedupCandidateIndexes,
+    DedupResolutionState,
+    _build_candidate_indexes,
+    _resolve_with_similarity,
+)
+from graphiti_core.utils.maintenance.edge_operations import (
+    filter_existing_duplicate_of_edges,
+)
+from graphiti_core.utils.text_utils import MAX_SUMMARY_CHARS, truncate_at_sentence
 
 logger = logging.getLogger(__name__)
 
+NodeSummaryFilter = Callable[[EntityNode], Awaitable[bool]]
+
 
 async def extract_nodes_reflexion(
     llm_client: LLMClient,
     episode: EpisodicNode,
     previous_episodes: list[EpisodicNode],
     node_names: list[str],
+    group_id: str | None = None,
 ) -> list[str]:
     # Prepare context for LLM
     context = {
@@ -59,7 +75,10 @@ async def extract_nodes_reflexion(
     }
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.reflexion(context), MissedEntities
+        prompt_library.extract_nodes.reflexion(context),
+        MissedEntities,
+        group_id=group_id,
+        prompt_name='extract_nodes.reflexion',
     )
     missed_entities = llm_response.get('missed_entities', [])
 
@@ -70,7 +89,7 @@ async def extract_nodes(
     clients: GraphitiClients,
     episode: EpisodicNode,
     previous_episodes: list[EpisodicNode],
-    entity_types: dict[str, BaseModel] | None = None,
+    entity_types: dict[str, type[BaseModel]] | None = None,
     excluded_entity_types: list[str] | None = None,
 ) -> list[EntityNode]:
     start = time()
@@ -115,20 +134,27 @@ async def extract_nodes(
             llm_response = await llm_client.generate_response(
                 prompt_library.extract_nodes.extract_message(context),
                 response_model=ExtractedEntities,
+                group_id=episode.group_id,
+                prompt_name='extract_nodes.extract_message',
             )
         elif episode.source == EpisodeType.text:
             llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_text(context), response_model=ExtractedEntities
+                prompt_library.extract_nodes.extract_text(context),
+                response_model=ExtractedEntities,
+                group_id=episode.group_id,
+                prompt_name='extract_nodes.extract_text',
             )
         elif episode.source == EpisodeType.json:
             llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_json(context), response_model=ExtractedEntities
+                prompt_library.extract_nodes.extract_json(context),
+                response_model=ExtractedEntities,
+                group_id=episode.group_id,
+                prompt_name='extract_nodes.extract_json',
             )
 
-        extracted_entities: list[ExtractedEntity] = [
-            ExtractedEntity(**entity_types_context)
-            for entity_types_context in llm_response.get('extracted_entities', [])
-        ]
+        response_object = ExtractedEntities(**llm_response)
+
+        extracted_entities: list[ExtractedEntity] = response_object.extracted_entities
 
         reflexion_iterations += 1
         if reflexion_iterations < MAX_REFLEXION_ITERATIONS:
@@ -137,6 +163,7 @@ async def extract_nodes(
                 episode,
                 previous_episodes,
                 [entity.name for entity in extracted_entities],
+                episode.group_id,
             )
 
             entities_missed = len(missing_entities) != 0
@@ -151,9 +178,13 @@ async def extract_nodes(
     # Convert the extracted data into EntityNode objects
     extracted_nodes = []
     for extracted_entity in filtered_extracted_entities:
-        entity_type_name = entity_types_context[extracted_entity.entity_type_id].get(
-            'entity_type_name'
-        )
+        type_id = extracted_entity.entity_type_id
+        if 0 <= type_id < len(entity_types_context):
+            entity_type_name = entity_types_context[extracted_entity.entity_type_id].get(
+                'entity_type_name'
+            )
+        else:
+            entity_type_name = 'Entity'
 
         # Check if this entity type should be excluded
         if excluded_entity_types and entity_type_name in excluded_entity_types:
@@ -173,20 +204,16 @@ async def extract_nodes(
         logger.debug(f'Created new node: {new_node.name} (UUID: {new_node.uuid})')
 
     logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
+
     return extracted_nodes
 
 
-async def resolve_extracted_nodes(
+async def _collect_candidate_nodes(
     clients: GraphitiClients,
     extracted_nodes: list[EntityNode],
-    episode: EpisodicNode | None = None,
-    previous_episodes: list[EpisodicNode] | None = None,
-    entity_types: dict[str, BaseModel] | None = None,
-    existing_nodes_override: list[EntityNode] | None = None,
-) -> tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]:
-    llm_client = clients.llm_client
-    driver = clients.driver
-
+    existing_nodes_override: list[EntityNode] | None,
+) -> list[EntityNode]:
+    """Search per extracted name and return unique candidates with overrides honored in order."""
     search_results: list[SearchResults] = await semaphore_gather(
         *[
             search(
@@ -200,33 +227,43 @@ async def resolve_extracted_nodes(
         ]
     )
 
-    candidate_nodes: list[EntityNode] = (
-        [node for result in search_results for node in result.nodes]
-        if existing_nodes_override is None
-        else existing_nodes_override
-    )
+    candidate_nodes: list[EntityNode] = [node for result in search_results for node in result.nodes]
 
-    existing_nodes_dict: dict[str, EntityNode] = {node.uuid: node for node in candidate_nodes}
+    if existing_nodes_override is not None:
+        candidate_nodes.extend(existing_nodes_override)
 
-    existing_nodes: list[EntityNode] = list(existing_nodes_dict.values())
+    seen_candidate_uuids: set[str] = set()
+    ordered_candidates: list[EntityNode] = []
+    for candidate in candidate_nodes:
+        if candidate.uuid in seen_candidate_uuids:
+            continue
+        seen_candidate_uuids.add(candidate.uuid)
+        ordered_candidates.append(candidate)
 
-    existing_nodes_context = (
-        [
-            {
-                **{
-                    'idx': i,
-                    'name': candidate.name,
-                    'entity_types': candidate.labels,
-                },
-                **candidate.attributes,
-            }
-            for i, candidate in enumerate(existing_nodes)
-        ],
-    )
+    return ordered_candidates
 
-    entity_types_dict: dict[str, BaseModel] = entity_types if entity_types is not None else {}
 
-    # Prepare context for LLM
+async def _resolve_with_llm(
+    llm_client: LLMClient,
+    extracted_nodes: list[EntityNode],
+    indexes: DedupCandidateIndexes,
+    state: DedupResolutionState,
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+    entity_types: dict[str, type[BaseModel]] | None,
+) -> None:
+    """Escalate unresolved nodes to the dedupe prompt so the LLM can select or reject duplicates.
+
+    The guardrails below defensively ignore malformed or duplicate LLM responses so the
+    ingestion workflow remains deterministic even when the model misbehaves.
+    """
+    if not state.unresolved_indices:
+        return
+
+    entity_types_dict: dict[str, type[BaseModel]] = entity_types if entity_types is not None else {}
+
+    llm_extracted_nodes = [extracted_nodes[i] for i in state.unresolved_indices]
+
     extracted_nodes_context = [
         {
             'id': i,
@@ -237,60 +274,180 @@ async def resolve_extracted_nodes(
             ).__doc__
             or 'Default Entity Type',
         }
-        for i, node in enumerate(extracted_nodes)
+        for i, node in enumerate(llm_extracted_nodes)
+    ]
+
+    sent_ids = [ctx['id'] for ctx in extracted_nodes_context]
+    logger.debug(
+        'Sending %d entities to LLM for deduplication with IDs 0-%d (actual IDs sent: %s)',
+        len(llm_extracted_nodes),
+        len(llm_extracted_nodes) - 1,
+        sent_ids if len(sent_ids) < 20 else f'{sent_ids[:10]}...{sent_ids[-10:]}',
+    )
+    if llm_extracted_nodes:
+        sample_size = min(3, len(extracted_nodes_context))
+        logger.debug(
+            'First %d entities: %s',
+            sample_size,
+            [(ctx['id'], ctx['name']) for ctx in extracted_nodes_context[:sample_size]],
+        )
+        if len(extracted_nodes_context) > 3:
+            logger.debug(
+                'Last %d entities: %s',
+                sample_size,
+                [(ctx['id'], ctx['name']) for ctx in extracted_nodes_context[-sample_size:]],
+            )
+
+    existing_nodes_context = [
+        {
+            **{
+                'idx': i,
+                'name': candidate.name,
+                'entity_types': candidate.labels,
+            },
+            **candidate.attributes,
+        }
+        for i, candidate in enumerate(indexes.existing_nodes)
     ]
 
     context = {
         'extracted_nodes': extracted_nodes_context,
         'existing_nodes': existing_nodes_context,
         'episode_content': episode.content if episode is not None else '',
-        'previous_episodes': [ep.content for ep in previous_episodes]
-        if previous_episodes is not None
-        else [],
+        'previous_episodes': (
+            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
+        ),
     }
 
     llm_response = await llm_client.generate_response(
         prompt_library.dedupe_nodes.nodes(context),
         response_model=NodeResolutions,
+        prompt_name='dedupe_nodes.nodes',
     )
 
-    node_resolutions: list = llm_response.get('entity_resolutions', [])
+    node_resolutions: list[NodeDuplicate] = NodeResolutions(**llm_response).entity_resolutions
 
-    resolved_nodes: list[EntityNode] = []
-    uuid_map: dict[str, str] = {}
-    node_duplicates: list[tuple[EntityNode, EntityNode]] = []
-    for resolution in node_resolutions:
-        resolution_id: int = resolution.get('id', -1)
-        duplicate_idx: int = resolution.get('duplicate_idx', -1)
+    valid_relative_range = range(len(state.unresolved_indices))
+    processed_relative_ids: set[int] = set()
 
-        extracted_node = extracted_nodes[resolution_id]
+    received_ids = {r.id for r in node_resolutions}
+    expected_ids = set(valid_relative_range)
+    missing_ids = expected_ids - received_ids
+    extra_ids = received_ids - expected_ids
+
+    logger.debug(
+        'Received %d resolutions for %d entities',
+        len(node_resolutions),
+        len(state.unresolved_indices),
+    )
 
-        resolved_node = (
-            existing_nodes[duplicate_idx]
-            if 0 <= duplicate_idx < len(existing_nodes)
-            else extracted_node
+    if missing_ids:
+        logger.warning('LLM did not return resolutions for IDs: %s', sorted(missing_ids))
+
+    if extra_ids:
+        logger.warning(
+            'LLM returned invalid IDs outside valid range 0-%d: %s (all returned IDs: %s)',
+            len(state.unresolved_indices) - 1,
+            sorted(extra_ids),
+            sorted(received_ids),
         )
 
-        # resolved_node.name = resolution.get('name')
+    for resolution in node_resolutions:
+        relative_id: int = resolution.id
+        duplicate_idx: int = resolution.duplicate_idx
+
+        if relative_id not in valid_relative_range:
+            logger.warning(
+                'Skipping invalid LLM dedupe id %d (valid range: 0-%d, received %d resolutions)',
+                relative_id,
+                len(state.unresolved_indices) - 1,
+                len(node_resolutions),
+            )
+            continue
 
-        resolved_nodes.append(resolved_node)
-        uuid_map[extracted_node.uuid] = resolved_node.uuid
+        if relative_id in processed_relative_ids:
+            logger.warning('Duplicate LLM dedupe id %s received; ignoring.', relative_id)
+            continue
+        processed_relative_ids.add(relative_id)
+
+        original_index = state.unresolved_indices[relative_id]
+        extracted_node = extracted_nodes[original_index]
+
+        resolved_node: EntityNode
+        if duplicate_idx == -1:
+            resolved_node = extracted_node
+        elif 0 <= duplicate_idx < len(indexes.existing_nodes):
+            resolved_node = indexes.existing_nodes[duplicate_idx]
+        else:
+            logger.warning(
+                'Invalid duplicate_idx %s for extracted node %s; treating as no duplicate.',
+                duplicate_idx,
+                extracted_node.uuid,
+            )
+            resolved_node = extracted_node
 
-        duplicates: list[int] = resolution.get('duplicates', [])
-        if duplicate_idx not in duplicates and duplicate_idx > -1:
-            duplicates.append(duplicate_idx)
-        for idx in duplicates:
-            existing_node = existing_nodes[idx] if idx < len(existing_nodes) else resolved_node
+        state.resolved_nodes[original_index] = resolved_node
+        state.uuid_map[extracted_node.uuid] = resolved_node.uuid
+        if resolved_node.uuid != extracted_node.uuid:
+            state.duplicate_pairs.append((extracted_node, resolved_node))
 
-            node_duplicates.append((extracted_node, existing_node))
 
-    logger.debug(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
+async def resolve_extracted_nodes(
+    clients: GraphitiClients,
+    extracted_nodes: list[EntityNode],
+    episode: EpisodicNode | None = None,
+    previous_episodes: list[EpisodicNode] | None = None,
+    entity_types: dict[str, type[BaseModel]] | None = None,
+    existing_nodes_override: list[EntityNode] | None = None,
+) -> tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]:
+    """Search for existing nodes, resolve deterministic matches, then escalate holdouts to the LLM dedupe prompt."""
+    llm_client = clients.llm_client
+    driver = clients.driver
+    existing_nodes = await _collect_candidate_nodes(
+        clients,
+        extracted_nodes,
+        existing_nodes_override,
+    )
+
+    indexes: DedupCandidateIndexes = _build_candidate_indexes(existing_nodes)
+
+    state = DedupResolutionState(
+        resolved_nodes=[None] * len(extracted_nodes),
+        uuid_map={},
+        unresolved_indices=[],
+    )
+
+    _resolve_with_similarity(extracted_nodes, indexes, state)
+
+    await _resolve_with_llm(
+        llm_client,
+        extracted_nodes,
+        indexes,
+        state,
+        episode,
+        previous_episodes,
+        entity_types,
+    )
+
+    for idx, node in enumerate(extracted_nodes):
+        if state.resolved_nodes[idx] is None:
+            state.resolved_nodes[idx] = node
+            state.uuid_map[node.uuid] = node.uuid
+
+    logger.debug(
+        'Resolved nodes: %s',
+        [(node.name, node.uuid) for node in state.resolved_nodes if node is not None],
+    )
 
     new_node_duplicates: list[
         tuple[EntityNode, EntityNode]
-    ] = await filter_existing_duplicate_of_edges(driver, node_duplicates)
+    ] = await filter_existing_duplicate_of_edges(driver, state.duplicate_pairs)
 
-    return resolved_nodes, uuid_map, new_node_duplicates
+    return (
+        [node for node in state.resolved_nodes if node is not None],
+        state.uuid_map,
+        new_node_duplicates,
+    )
 
 
 async def extract_attributes_from_nodes(
@@ -298,7 +455,8 @@ async def extract_attributes_from_nodes(
     nodes: list[EntityNode],
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
-    entity_types: dict[str, BaseModel] | None = None,
+    entity_types: dict[str, type[BaseModel]] | None = None,
+    should_summarize_node: NodeSummaryFilter | None = None,
 ) -> list[EntityNode]:
     llm_client = clients.llm_client
     embedder = clients.embedder
@@ -309,9 +467,12 @@ async def extract_attributes_from_nodes(
                 node,
                 episode,
                 previous_episodes,
-                entity_types.get(next((item for item in node.labels if item != 'Entity'), ''))
-                if entity_types is not None
-                else None,
+                (
+                    entity_types.get(next((item for item in node.labels if item != 'Entity'), ''))
+                    if entity_types is not None
+                    else None
+                ),
+                should_summarize_node,
             )
             for node in nodes
         ]
@@ -327,99 +488,100 @@ async def extract_attributes_from_node(
     node: EntityNode,
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
-    entity_type: BaseModel | None = None,
+    entity_type: type[BaseModel] | None = None,
+    should_summarize_node: NodeSummaryFilter | None = None,
 ) -> EntityNode:
-    node_context: dict[str, Any] = {
-        'name': node.name,
-        'summary': node.summary,
-        'entity_types': node.labels,
-        'attributes': node.attributes,
-    }
+    # Extract attributes if entity type is defined and has attributes
+    llm_response = await _extract_entity_attributes(
+        llm_client, node, episode, previous_episodes, entity_type
+    )
 
-    attributes_definitions: dict[str, Any] = {
-        'summary': (
-            str,
-            Field(
-                description='Summary containing the important information about the entity. Under 250 words',
-            ),
-        )
-    }
+    # Extract summary if needed
+    await _extract_entity_summary(
+        llm_client, node, episode, previous_episodes, should_summarize_node
+    )
 
-    if entity_type is not None:
-        for field_name, field_info in entity_type.model_fields.items():
-            attributes_definitions[field_name] = (
-                field_info.annotation,
-                Field(description=field_info.description),
-            )
+    node.attributes.update(llm_response)
 
-    unique_model_name = f'EntityAttributes_{uuid4().hex}'
-    entity_attributes_model = pydantic.create_model(unique_model_name, **attributes_definitions)
+    return node
 
-    summary_context: dict[str, Any] = {
-        'node': node_context,
-        'episode_content': episode.content if episode is not None else '',
-        'previous_episodes': [ep.content for ep in previous_episodes]
-        if previous_episodes is not None
-        else [],
-    }
+
+async def _extract_entity_attributes(
+    llm_client: LLMClient,
+    node: EntityNode,
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+    entity_type: type[BaseModel] | None,
+) -> dict[str, Any]:
+    if entity_type is None or len(entity_type.model_fields) == 0:
+        return {}
+
+    attributes_context = _build_episode_context(
+        # should not include summary
+        node_data={
+            'name': node.name,
+            'entity_types': node.labels,
+            'attributes': node.attributes,
+        },
+        episode=episode,
+        previous_episodes=previous_episodes,
+    )
 
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.extract_attributes(summary_context),
-        response_model=entity_attributes_model,
+        prompt_library.extract_nodes.extract_attributes(attributes_context),
+        response_model=entity_type,
         model_size=ModelSize.small,
+        group_id=node.group_id,
+        prompt_name='extract_nodes.extract_attributes',
     )
 
-    node.summary = llm_response.get('summary', node.summary)
-    node_attributes = {key: value for key, value in llm_response.items()}
-
-    with suppress(KeyError):
-        del node_attributes['summary']
+    # validate response
+    entity_type(**llm_response)
 
-    node.attributes.update(node_attributes)
+    return llm_response
 
-    return node
 
-
-async def dedupe_node_list(
+async def _extract_entity_summary(
     llm_client: LLMClient,
-    nodes: list[EntityNode],
-) -> tuple[list[EntityNode], dict[str, str]]:
-    start = time()
-
-    # build node map
-    node_map = {}
-    for node in nodes:
-        node_map[node.uuid] = node
-
-    # Prepare context for LLM
-    nodes_context = [{'uuid': node.uuid, 'name': node.name, **node.attributes} for node in nodes]
-
-    context = {
-        'nodes': nodes_context,
-    }
-
-    llm_response = await llm_client.generate_response(
-        prompt_library.dedupe_nodes.node_list(context)
+    node: EntityNode,
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+    should_summarize_node: NodeSummaryFilter | None,
+) -> None:
+    if should_summarize_node is not None and not await should_summarize_node(node):
+        return
+
+    summary_context = _build_episode_context(
+        node_data={
+            'name': node.name,
+            'summary': truncate_at_sentence(node.summary, MAX_SUMMARY_CHARS),
+            'entity_types': node.labels,
+            'attributes': node.attributes,
+        },
+        episode=episode,
+        previous_episodes=previous_episodes,
     )
 
-    nodes_data = llm_response.get('nodes', [])
+    summary_response = await llm_client.generate_response(
+        prompt_library.extract_nodes.extract_summary(summary_context),
+        response_model=EntitySummary,
+        model_size=ModelSize.small,
+        group_id=node.group_id,
+        prompt_name='extract_nodes.extract_summary',
+    )
 
-    end = time()
-    logger.debug(f'Deduplicated nodes: {nodes_data} in {(end - start) * 1000} ms')
-
-    # Get full node data
-    unique_nodes = []
-    uuid_map: dict[str, str] = {}
-    for node_data in nodes_data:
-        node_instance: EntityNode | None = node_map.get(node_data['uuids'][0])
-        if node_instance is None:
-            logger.warning(f'Node {node_data["uuids"][0]} not found in node map')
-            continue
-        node_instance.summary = node_data['summary']
-        unique_nodes.append(node_instance)
+    node.summary = truncate_at_sentence(summary_response.get('summary', ''), MAX_SUMMARY_CHARS)
 
-        for uuid in node_data['uuids'][1:]:
-            uuid_value = node_map[node_data['uuids'][0]].uuid
-            uuid_map[uuid] = uuid_value
 
-    return unique_nodes, uuid_map
+def _build_episode_context(
+    node_data: dict[str, Any],
+    episode: EpisodicNode | None,
+    previous_episodes: list[EpisodicNode] | None,
+) -> dict[str, Any]:
+    return {
+        'node': node_data,
+        'episode_content': episode.content if episode is not None else '',
+        'previous_episodes': (
+            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
+        ),
+    }

graphiti_core/utils/maintenance/temporal_operations.py

@@ -43,7 +43,9 @@ async def extract_edge_dates(
         'reference_timestamp': current_episode.valid_at.isoformat(),
     }
     llm_response = await llm_client.generate_response(
-        prompt_library.extract_edge_dates.v1(context), response_model=EdgeDates
+        prompt_library.extract_edge_dates.v1(context),
+        response_model=EdgeDates,
+        prompt_name='extract_edge_dates.v1',
     )
 
     valid_at = llm_response.get('valid_at')
@@ -70,7 +72,9 @@ async def extract_edge_dates(
 
 
 async def get_edge_contradictions(
-    llm_client: LLMClient, new_edge: EntityEdge, existing_edges: list[EntityEdge]
+    llm_client: LLMClient,
+    new_edge: EntityEdge,
+    existing_edges: list[EntityEdge],
 ) -> list[EntityEdge]:
     start = time()
 
@@ -79,12 +83,16 @@ async def get_edge_contradictions(
         {'id': i, 'fact': existing_edge.fact} for i, existing_edge in enumerate(existing_edges)
     ]
 
-    context = {'new_edge': new_edge_context, 'existing_edges': existing_edge_context}
+    context = {
+        'new_edge': new_edge_context,
+        'existing_edges': existing_edge_context,
+    }
 
     llm_response = await llm_client.generate_response(
         prompt_library.invalidate_edges.v2(context),
         response_model=InvalidatedEdges,
         model_size=ModelSize.small,
+        prompt_name='invalidate_edges.v2',
     )
 
     contradicted_facts: list[int] = llm_response.get('contradicted_facts', [])