graphiti-core 0.12.0rc1__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 NodeDuplicate, NodeResolutions
 from graphiti_core.prompts.extract_nodes import (
+    EntitySummary,
     ExtractedEntities,
     ExtractedEntity,
     MissedEntities,
@@ -40,15 +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.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 = {
@@ -58,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', [])
 
@@ -69,7 +89,8 @@ 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()
     llm_client = clients.llm_client
@@ -113,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:
@@ -135,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
@@ -149,9 +178,18 @@ 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:
+            logger.debug(f'Excluding entity "{extracted_entity.name}" of type "{entity_type_name}"')
+            continue
 
         labels: list[str] = list({'Entity', str(entity_type_name)})
 
@@ -166,68 +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 dedupe_extracted_nodes(
-    llm_client: LLMClient,
-    extracted_nodes: list[EntityNode],
-    existing_nodes: list[EntityNode],
-) -> tuple[list[EntityNode], dict[str, str]]:
-    start = time()
-
-    # build existing node map
-    node_map: dict[str, EntityNode] = {}
-    for node in existing_nodes:
-        node_map[node.uuid] = node
-
-    # Prepare context for LLM
-    existing_nodes_context = [
-        {'uuid': node.uuid, 'name': node.name, 'summary': node.summary} for node in existing_nodes
-    ]
-
-    extracted_nodes_context = [
-        {'uuid': node.uuid, 'name': node.name, 'summary': node.summary} for node in extracted_nodes
-    ]
-
-    context = {
-        'existing_nodes': existing_nodes_context,
-        'extracted_nodes': extracted_nodes_context,
-    }
-
-    llm_response = await llm_client.generate_response(prompt_library.dedupe_nodes.node(context))
-
-    duplicate_data = llm_response.get('duplicates', [])
-
-    end = time()
-    logger.debug(f'Deduplicated nodes: {duplicate_data} in {(end - start) * 1000} ms')
-
-    uuid_map: dict[str, str] = {}
-    for duplicate in duplicate_data:
-        uuid_value = duplicate['duplicate_of']
-        uuid_map[duplicate['uuid']] = uuid_value
-
-    nodes: list[EntityNode] = []
-    for node in extracted_nodes:
-        if node.uuid in uuid_map:
-            existing_uuid = uuid_map[node.uuid]
-            existing_node = node_map[existing_uuid]
-            nodes.append(existing_node)
-        else:
-            nodes.append(node)
 
-    return nodes, uuid_map
+    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,
-) -> tuple[list[EntityNode], dict[str, str]]:
-    llm_client = clients.llm_client
-
+    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(
@@ -241,11 +227,43 @@ async def resolve_extracted_nodes(
         ]
     )
 
-    existing_nodes_lists: list[list[EntityNode]] = [result.nodes for result in search_results]
+    candidate_nodes: list[EntityNode] = [node for result in search_results for node in result.nodes]
 
-    entity_types_dict: dict[str, BaseModel] = entity_types if entity_types is not None else {}
+    if existing_nodes_override is not None:
+        candidate_nodes.extend(existing_nodes_override)
+
+    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)
+
+    return ordered_candidates
+
+
+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]
 
-    # Prepare context for LLM
     extracted_nodes_context = [
         {
             'id': i,
@@ -255,122 +273,181 @@ async def resolve_extracted_nodes(
                 next((item for item in node.labels if item != 'Entity'), '')
             ).__doc__
             or 'Default Entity Type',
-            'duplication_candidates': [
-                {
-                    **{
-                        'idx': j,
-                        'name': candidate.name,
-                        'entity_types': candidate.labels,
-                    },
-                    **candidate.attributes,
-                }
-                for j, candidate in enumerate(existing_nodes_lists[i])
-            ],
         }
-        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] = {}
-    for resolution in node_resolutions:
-        resolution_id = resolution.get('id', -1)
-        duplicate_idx = resolution.get('duplicate_idx', -1)
+    valid_relative_range = range(len(state.unresolved_indices))
+    processed_relative_ids: set[int] = set()
+
+    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),
+    )
 
-        extracted_node = extracted_nodes[resolution_id]
+    if missing_ids:
+        logger.warning('LLM did not return resolutions for IDs: %s', sorted(missing_ids))
 
-        resolved_node = (
-            existing_nodes_lists[resolution_id][duplicate_idx]
-            if 0 <= duplicate_idx < len(existing_nodes_lists[resolution_id])
-            else extracted_node
+    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)
 
-    logger.debug(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
+        original_index = state.unresolved_indices[relative_id]
+        extracted_node = extracted_nodes[original_index]
 
-    return resolved_nodes, uuid_map
+        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
 
+        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))
 
-async def resolve_extracted_node(
-    llm_client: LLMClient,
-    extracted_node: EntityNode,
-    existing_nodes: list[EntityNode],
+
+async def resolve_extracted_nodes(
+    clients: GraphitiClients,
+    extracted_nodes: list[EntityNode],
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
-    entity_type: BaseModel | None = None,
-) -> EntityNode:
-    start = time()
-    if len(existing_nodes) == 0:
-        return extracted_node
-
-    # Prepare context for LLM
-    existing_nodes_context = [
-        {
-            **{
-                'id': i,
-                'name': node.name,
-                'entity_types': node.labels,
-            },
-            **node.attributes,
-        }
-        for i, node in enumerate(existing_nodes)
-    ]
-
-    extracted_node_context = {
-        'name': extracted_node.name,
-        'entity_type': entity_type.__name__ if entity_type is not None else 'Entity',  # type: ignore
-    }
+    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,
+    )
 
-    context = {
-        'existing_nodes': existing_nodes_context,
-        'extracted_node': extracted_node_context,
-        'entity_type_description': entity_type.__doc__
-        if entity_type is not None
-        else 'Default Entity Type',
-        '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 [],
-    }
+    indexes: DedupCandidateIndexes = _build_candidate_indexes(existing_nodes)
 
-    llm_response = await llm_client.generate_response(
-        prompt_library.dedupe_nodes.node(context),
-        response_model=NodeDuplicate,
-        model_size=ModelSize.small,
+    state = DedupResolutionState(
+        resolved_nodes=[None] * len(extracted_nodes),
+        uuid_map={},
+        unresolved_indices=[],
     )
 
-    duplicate_id: int = llm_response.get('duplicate_node_id', -1)
+    _resolve_with_similarity(extracted_nodes, indexes, state)
 
-    node = (
-        existing_nodes[duplicate_id] if 0 <= duplicate_id < len(existing_nodes) else extracted_node
+    await _resolve_with_llm(
+        llm_client,
+        extracted_nodes,
+        indexes,
+        state,
+        episode,
+        previous_episodes,
+        entity_types,
     )
 
-    node.name = llm_response.get('name', '')
+    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
 
-    end = time()
     logger.debug(
-        f'Resolved node: {extracted_node.name} is {node.name}, in {(end - start) * 1000} ms'
+        'Resolved nodes: %s',
+        [(node.name, node.uuid) for node in state.resolved_nodes if node is not None],
     )
 
-    return node
+    new_node_duplicates: list[
+        tuple[EntityNode, EntityNode]
+    ] = await filter_existing_duplicate_of_edges(driver, state.duplicate_pairs)
+
+    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(
@@ -378,11 +455,11 @@ 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
-
     updated_nodes: list[EntityNode] = await semaphore_gather(
         *[
             extract_attributes_from_node(
@@ -390,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
         ]
@@ -408,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 node
+    return llm_response
 
 
-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 []
+        ),
+    }