graphiti-core 0.17.4__py3-none-any.whl → 0.25.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.helpers import 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,
@@ -39,17 +43,35 @@ from graphiti_core.search.search import search
 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.content_chunking import (
+    chunk_json_content,
+    chunk_message_content,
+    chunk_text_content,
+    should_chunk,
+)
 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 +81,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,26 +95,69 @@ 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,
+    custom_extraction_instructions: str | None = None,
 ) -> list[EntityNode]:
+    """Extract entity nodes from an episode with adaptive chunking.
+
+    For high-density content (many entities per token), the content is chunked
+    and processed in parallel to avoid LLM timeouts and truncation issues.
+    """
     start = time()
     llm_client = clients.llm_client
-    llm_response = {}
-    custom_prompt = ''
-    entities_missed = True
-    reflexion_iterations = 0
 
+    # Build entity types context
+    entity_types_context = _build_entity_types_context(entity_types)
+
+    # Build base context
+    context = {
+        'episode_content': episode.content,
+        'episode_timestamp': episode.valid_at.isoformat(),
+        'previous_episodes': [ep.content for ep in previous_episodes],
+        'custom_extraction_instructions': custom_extraction_instructions or '',
+        'entity_types': entity_types_context,
+        'source_description': episode.source_description,
+    }
+
+    # Check if chunking is needed (based on entity density)
+    if should_chunk(episode.content, episode.source):
+        extracted_entities = await _extract_nodes_chunked(llm_client, episode, context)
+    else:
+        extracted_entities = await _extract_nodes_single(llm_client, episode, context)
+
+    # Filter empty names
+    filtered_entities = [e for e in extracted_entities if e.name.strip()]
+
+    end = time()
+    logger.debug(f'Extracted {len(filtered_entities)} entities in {(end - start) * 1000:.0f} ms')
+
+    # Convert to EntityNode objects
+    extracted_nodes = _create_entity_nodes(
+        filtered_entities, entity_types_context, excluded_entity_types, episode
+    )
+
+    logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
+    return extracted_nodes
+
+
+def _build_entity_types_context(
+    entity_types: dict[str, type[BaseModel]] | None,
+) -> list[dict]:
+    """Build entity types context with ID mappings."""
     entity_types_context = [
         {
             'entity_type_id': 0,
             'entity_type_name': 'Entity',
-            'entity_type_description': 'Default entity classification. Use this entity type if the entity is not one of the other listed types.',
+            'entity_type_description': (
+                'Default entity classification. Use this entity type '
+                'if the entity is not one of the other listed types.'
+            ),
         }
     ]
 
-    entity_types_context += (
-        [
+    if entity_types is not None:
+        entity_types_context += [
             {
                 'entity_type_id': i + 1,
                 'entity_type_name': type_name,
@@ -97,63 +165,126 @@ async def extract_nodes(
             }
             for i, (type_name, type_model) in enumerate(entity_types.items())
         ]
-        if entity_types is not None
-        else []
+
+    return entity_types_context
+
+
+async def _extract_nodes_single(
+    llm_client: LLMClient,
+    episode: EpisodicNode,
+    context: dict,
+) -> list[ExtractedEntity]:
+    """Extract entities using a single LLM call."""
+    llm_response = await _call_extraction_llm(llm_client, episode, context)
+    response_object = ExtractedEntities(**llm_response)
+    return response_object.extracted_entities
+
+
+async def _extract_nodes_chunked(
+    llm_client: LLMClient,
+    episode: EpisodicNode,
+    context: dict,
+) -> list[ExtractedEntity]:
+    """Extract entities from large content using chunking."""
+    # Chunk the content based on episode type
+    if episode.source == EpisodeType.json:
+        chunks = chunk_json_content(episode.content)
+    elif episode.source == EpisodeType.message:
+        chunks = chunk_message_content(episode.content)
+    else:
+        chunks = chunk_text_content(episode.content)
+
+    logger.debug(f'Chunked content into {len(chunks)} chunks for entity extraction')
+
+    # Extract entities from each chunk in parallel
+    chunk_results = await semaphore_gather(
+        *[_extract_from_chunk(llm_client, chunk, context, episode) for chunk in chunks]
     )
 
-    context = {
-        'episode_content': episode.content,
-        'episode_timestamp': episode.valid_at.isoformat(),
-        'previous_episodes': [ep.content for ep in previous_episodes],
-        'custom_prompt': custom_prompt,
-        'entity_types': entity_types_context,
-        'source_description': episode.source_description,
-    }
+    # Merge and deduplicate entities across chunks
+    merged_entities = _merge_extracted_entities(chunk_results)
+    logger.debug(
+        f'Merged {sum(len(r) for r in chunk_results)} entities into {len(merged_entities)} unique'
+    )
 
-    while entities_missed and reflexion_iterations <= MAX_REFLEXION_ITERATIONS:
-        if episode.source == EpisodeType.message:
-            llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_message(context),
-                response_model=ExtractedEntities,
-            )
-        elif episode.source == EpisodeType.text:
-            llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_text(context), response_model=ExtractedEntities
-            )
-        elif episode.source == EpisodeType.json:
-            llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_json(context), response_model=ExtractedEntities
-            )
+    return merged_entities
 
-        extracted_entities: list[ExtractedEntity] = [
-            ExtractedEntity(**entity_types_context)
-            for entity_types_context in llm_response.get('extracted_entities', [])
-        ]
 
-        reflexion_iterations += 1
-        if reflexion_iterations < MAX_REFLEXION_ITERATIONS:
-            missing_entities = await extract_nodes_reflexion(
-                llm_client,
-                episode,
-                previous_episodes,
-                [entity.name for entity in extracted_entities],
-            )
+async def _extract_from_chunk(
+    llm_client: LLMClient,
+    chunk: str,
+    base_context: dict,
+    episode: EpisodicNode,
+) -> list[ExtractedEntity]:
+    """Extract entities from a single chunk."""
+    chunk_context = {**base_context, 'episode_content': chunk}
+    llm_response = await _call_extraction_llm(llm_client, episode, chunk_context)
+    return ExtractedEntities(**llm_response).extracted_entities
+
+
+async def _call_extraction_llm(
+    llm_client: LLMClient,
+    episode: EpisodicNode,
+    context: dict,
+) -> dict:
+    """Call the appropriate extraction prompt based on episode type."""
+    if episode.source == EpisodeType.message:
+        prompt = prompt_library.extract_nodes.extract_message(context)
+        prompt_name = 'extract_nodes.extract_message'
+    elif episode.source == EpisodeType.text:
+        prompt = prompt_library.extract_nodes.extract_text(context)
+        prompt_name = 'extract_nodes.extract_text'
+    elif episode.source == EpisodeType.json:
+        prompt = prompt_library.extract_nodes.extract_json(context)
+        prompt_name = 'extract_nodes.extract_json'
+    else:
+        # Fallback to text extraction
+        prompt = prompt_library.extract_nodes.extract_text(context)
+        prompt_name = 'extract_nodes.extract_text'
+
+    return await llm_client.generate_response(
+        prompt,
+        response_model=ExtractedEntities,
+        group_id=episode.group_id,
+        prompt_name=prompt_name,
+    )
 
-            entities_missed = len(missing_entities) != 0
 
-            custom_prompt = 'Make sure that the following entities are extracted: '
-            for entity in missing_entities:
-                custom_prompt += f'\n{entity},'
+def _merge_extracted_entities(
+    chunk_results: list[list[ExtractedEntity]],
+) -> list[ExtractedEntity]:
+    """Merge entities from multiple chunks, deduplicating by normalized name.
 
-    filtered_extracted_entities = [entity for entity in extracted_entities if entity.name.strip()]
-    end = time()
-    logger.debug(f'Extracted new nodes: {filtered_extracted_entities} in {(end - start) * 1000} ms')
-    # Convert the extracted data into EntityNode objects
+    When duplicates occur, prefer the first occurrence (maintains ordering).
+    """
+    seen_names: set[str] = set()
+    merged: list[ExtractedEntity] = []
+
+    for entities in chunk_results:
+        for entity in entities:
+            normalized = entity.name.strip().lower()
+            if normalized and normalized not in seen_names:
+                seen_names.add(normalized)
+                merged.append(entity)
+
+    return merged
+
+
+def _create_entity_nodes(
+    extracted_entities: list[ExtractedEntity],
+    entity_types_context: list[dict],
+    excluded_entity_types: list[str] | None,
+    episode: EpisodicNode,
+) -> list[EntityNode]:
+    """Convert ExtractedEntity objects to 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'
-        )
+
+    for extracted_entity in extracted_entities:
+        type_id = extracted_entity.entity_type_id
+        if 0 <= type_id < len(entity_types_context):
+            entity_type_name = entity_types_context[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:
@@ -172,21 +303,15 @@ async def extract_nodes(
         extracted_nodes.append(new_node)
         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 +325,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 +372,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()
+
+    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
 
-        extracted_node = extracted_nodes[resolution_id]
+    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
+
+        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
+
+        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_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,
+    )
 
-        resolved_nodes.append(resolved_node)
-        uuid_map[extracted_node.uuid] = resolved_node.uuid
+    indexes: DedupCandidateIndexes = _build_candidate_indexes(existing_nodes)
 
-        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 = DedupResolutionState(
+        resolved_nodes=[None] * len(extracted_nodes),
+        uuid_map={},
+        unresolved_indices=[],
+    )
 
-            node_duplicates.append((extracted_node, existing_node))
+    _resolve_with_similarity(extracted_nodes, indexes, state)
 
-    logger.debug(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
+    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 +553,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 +565,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 +586,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 []
+        ),
+    }