getzep--graphiti/tests/utils/maintenance/test_node_operations.py

import logging
from collections import defaultdict
from unittest.mock import AsyncMock, MagicMock

import pytest

from graphiti_core.graphiti_types import GraphitiClients
from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
from graphiti_core.utils.datetime_utils import utc_now
from graphiti_core.utils.maintenance.dedup_helpers import (
    DedupCandidateIndexes,
    DedupResolutionState,
    _build_candidate_indexes,
    _cached_shingles,
    _has_high_entropy,
    _hash_shingle,
    _jaccard_similarity,
    _lsh_bands,
    _minhash_signature,
    _name_entropy,
    _normalize_name_for_fuzzy,
    _normalize_string_exact,
    _resolve_with_similarity,
    _shingles,
)
from graphiti_core.utils.maintenance.node_operations import (
    _collect_candidate_nodes,
    _extract_entity_summaries_batch,
    _resolve_with_llm,
    extract_attributes_from_nodes,
    resolve_extracted_nodes,
)


def _make_clients():
    driver = MagicMock()
    embedder = MagicMock()
    cross_encoder = MagicMock()
    llm_client = MagicMock()
    llm_generate = AsyncMock()
    llm_client.generate_response = llm_generate

    clients = GraphitiClients.model_construct(  # bypass validation to allow test doubles
        driver=driver,
        embedder=embedder,
        cross_encoder=cross_encoder,
        llm_client=llm_client,
    )

    return clients, llm_generate


def _make_episode(group_id: str = 'group'):
    return EpisodicNode(
        name='episode',
        group_id=group_id,
        source=EpisodeType.message,
        source_description='test',
        content='content',
        valid_at=utc_now(),
    )


def _semantic_candidates(candidate_groups: list[list[EntityNode]]):
    async def fake_search(*_, **__):
        return candidate_groups

    return fake_search


@pytest.mark.asyncio
async def test_resolve_nodes_exact_match_skips_llm(monkeypatch):
    clients, llm_generate = _make_clients()

    candidate = EntityNode(name='Joe Michaels', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Joe Michaels', group_id='group', labels=['Entity'])
    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[candidate]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == candidate.uuid
    assert uuid_map[extracted.uuid] == candidate.uuid
    llm_generate.assert_not_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_exact_match_promotes_generic_candidate_type(monkeypatch):
    clients, llm_generate = _make_clients()

    candidate = EntityNode(name='Audrey', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Audrey', group_id='group', labels=['Entity', 'Person'])
    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[candidate]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == candidate.uuid
    assert set(resolved[0].labels) == {'Entity', 'Person'}
    assert set(candidate.labels) == {'Entity', 'Person'}
    assert uuid_map[extracted.uuid] == candidate.uuid
    llm_generate.assert_not_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_low_entropy_uses_llm(monkeypatch):
    clients, llm_generate = _make_clients()
    llm_generate.return_value = {
        'entity_resolutions': [
            {
                'id': 0,
                'name': 'Joe',
                'duplicate_candidate_id': -1,
            }
        ]
    }

    candidate = EntityNode(name='Joseph', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Joe', group_id='group', labels=['Entity'])
    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[candidate]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == extracted.uuid
    assert uuid_map[extracted.uuid] == extracted.uuid
    llm_generate.assert_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_short_name_exact_match_skips_llm(monkeypatch):
    """Short names with a unique exact candidate match should resolve deterministically."""
    clients, llm_generate = _make_clients()

    candidate = EntityNode(name='Java', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Java', group_id='group', labels=['Entity'])
    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[candidate]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == candidate.uuid
    assert uuid_map[extracted.uuid] == candidate.uuid
    llm_generate.assert_not_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_fuzzy_match(monkeypatch):
    clients, llm_generate = _make_clients()

    candidate = EntityNode(name='Joe-Michaels', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Joe Michaels', group_id='group', labels=['Entity'])
    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[candidate]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == candidate.uuid
    assert uuid_map[extracted.uuid] == candidate.uuid
    llm_generate.assert_not_awaited()


@pytest.mark.asyncio
async def test_collect_candidate_nodes_dedupes_and_merges_override(monkeypatch):
    clients, _ = _make_clients()

    candidate = EntityNode(name='Alice', group_id='group', labels=['Entity'])
    override_duplicate = EntityNode(
        uuid=candidate.uuid,
        name='Alice Alt',
        group_id='group',
        labels=['Entity'],
    )
    extracted = EntityNode(name='Alice', group_id='group', labels=['Entity'])

    semantic_search_mock = AsyncMock(return_value=[[candidate]])
    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        semantic_search_mock,
    )

    result = await _collect_candidate_nodes(
        clients,
        [extracted],
        existing_nodes_override=[override_duplicate],
    )

    assert len(result) == 1
    assert len(result[0]) == 1
    assert result[0][0].uuid == candidate.uuid
    semantic_search_mock.assert_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_semantic_miss_keeps_node_without_llm(monkeypatch):
    clients, llm_generate = _make_clients()
    extracted = EntityNode(name='Completely New Thing', group_id='group', labels=['Entity'])

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[]]),
    )

    resolved, uuid_map, duplicates = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == extracted.uuid
    assert uuid_map[extracted.uuid] == extracted.uuid
    assert duplicates == []
    llm_generate.assert_not_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_multiple_exact_matches_use_llm(monkeypatch):
    clients, llm_generate = _make_clients()
    llm_generate.return_value = {
        'entity_resolutions': [
            {
                'id': 0,
                'name': 'Java',
                'duplicate_candidate_id': 0,
            }
        ]
    }
    candidate1 = EntityNode(name='Java', group_id='group', labels=['Entity'])
    candidate2 = EntityNode(name='Java', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Java', group_id='group', labels=['Entity'])

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[candidate1, candidate2]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        [extracted],
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == candidate1.uuid
    assert uuid_map[extracted.uuid] == candidate1.uuid
    llm_generate.assert_awaited()


@pytest.mark.asyncio
async def test_resolve_nodes_batches_unresolved_nodes_into_one_llm_call(monkeypatch):
    clients, llm_generate = _make_clients()
    llm_generate.return_value = {
        'entity_resolutions': [
            {
                'id': 0,
                'name': 'Joe',
                'duplicate_candidate_id': -1,
            },
            {
                'id': 1,
                'name': 'Java',
                'duplicate_candidate_id': 1,
            },
        ]
    }

    low_entropy_candidate = EntityNode(name='Joseph', group_id='group', labels=['Entity'])
    java_candidate_1 = EntityNode(name='Java', group_id='group', labels=['Entity'])
    java_candidate_2 = EntityNode(name='Java', group_id='group', labels=['Entity'])
    extracted_nodes = [
        EntityNode(name='Joe', group_id='group', labels=['Entity']),
        EntityNode(name='Java', group_id='group', labels=['Entity']),
    ]

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations._semantic_candidate_search',
        _semantic_candidates([[low_entropy_candidate], [java_candidate_1, java_candidate_2]]),
    )

    resolved, uuid_map, _ = await resolve_extracted_nodes(
        clients,
        extracted_nodes,
        episode=_make_episode(),
        previous_episodes=[],
    )

    assert resolved[0].uuid == extracted_nodes[0].uuid
    assert resolved[1].uuid == java_candidate_1.uuid
    assert uuid_map[extracted_nodes[0].uuid] == extracted_nodes[0].uuid
    assert uuid_map[extracted_nodes[1].uuid] == java_candidate_1.uuid
    assert llm_generate.await_count == 1


def test_build_candidate_indexes_populates_structures():
    candidate = EntityNode(name='Bob Dylan', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate])

    normalized_key = candidate.name.lower()
    assert indexes.normalized_existing[normalized_key][0].uuid == candidate.uuid
    assert indexes.nodes_by_uuid[candidate.uuid] is candidate
    assert candidate.uuid in indexes.shingles_by_candidate
    assert any(candidate.uuid in bucket for bucket in indexes.lsh_buckets.values())


def test_normalize_helpers():
    assert _normalize_string_exact('  Alice   Smith ') == 'alice smith'
    assert _normalize_name_for_fuzzy('Alice-Smith!') == 'alice smith'


def test_name_entropy_variants():
    assert _name_entropy('alice') > _name_entropy('aaaaa')
    assert _name_entropy('') == 0.0


def test_has_high_entropy_rules():
    assert _has_high_entropy('meaningful name') is True
    assert _has_high_entropy('aa') is False


def test_shingles_and_cache():
    raw = 'alice'
    shingle_set = _shingles(raw)
    assert shingle_set == {'ali', 'lic', 'ice'}
    assert _cached_shingles(raw) == shingle_set
    assert _cached_shingles(raw) is _cached_shingles(raw)


def test_hash_minhash_and_lsh():
    shingles = {'abc', 'bcd', 'cde'}
    signature = _minhash_signature(shingles)
    assert len(signature) == 32
    bands = _lsh_bands(signature)
    assert all(len(band) == 4 for band in bands)
    hashed = {_hash_shingle(s, 0) for s in shingles}
    assert len(hashed) == len(shingles)


def test_jaccard_similarity_edges():
    a = {'a', 'b'}
    b = {'a', 'c'}
    assert _jaccard_similarity(a, b) == pytest.approx(1 / 3)
    assert _jaccard_similarity(set(), set()) == 1.0
    assert _jaccard_similarity(a, set()) == 0.0


def test_resolve_with_similarity_exact_match_updates_state():
    candidate = EntityNode(name='Charlie Parker', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Charlie Parker', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[])

    _resolve_with_similarity([extracted], indexes, state)

    assert state.resolved_nodes[0].uuid == candidate.uuid
    assert state.uuid_map[extracted.uuid] == candidate.uuid
    assert state.unresolved_indices == []
    assert state.duplicate_pairs == [(extracted, candidate)]


def test_resolve_with_similarity_short_name_exact_match_resolves_deterministically():
    """Short names like 'Nate' should resolve via exact match without hitting the LLM."""
    candidate = EntityNode(name='Nate', group_id='group', labels=['Entity', 'Person'])
    extracted = EntityNode(name='Nate', group_id='group', labels=['Entity', 'Person'])

    indexes = _build_candidate_indexes([candidate])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[])

    _resolve_with_similarity([extracted], indexes, state)

    assert state.resolved_nodes[0].uuid == candidate.uuid
    assert state.uuid_map[extracted.uuid] == candidate.uuid
    assert state.unresolved_indices == []
    assert state.duplicate_pairs == [(extracted, candidate)]


def test_resolve_with_similarity_short_name_no_candidate_defers_to_llm():
    """Short names with no exact match should still reach the LLM for resolution."""
    extracted = EntityNode(name='Nate', group_id='group', labels=['Entity', 'Person'])

    indexes = _build_candidate_indexes([])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[])

    _resolve_with_similarity([extracted], indexes, state)

    assert state.resolved_nodes[0] is None
    assert state.uuid_map == {}
    assert state.unresolved_indices == [0]
    assert state.duplicate_pairs == []


def test_resolve_with_similarity_short_name_multiple_candidates_defers_to_llm():
    """Short names with multiple exact matches should escalate to LLM."""
    candidate1 = EntityNode(name='Java', group_id='group', labels=['Entity'])
    candidate2 = EntityNode(name='Java', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Java', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate1, candidate2])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[])

    _resolve_with_similarity([extracted], indexes, state)

    assert state.resolved_nodes[0] is None
    assert state.uuid_map == {}
    assert state.unresolved_indices == [0]
    assert state.duplicate_pairs == []


def test_resolve_with_similarity_low_entropy_defers_resolution():
    extracted = EntityNode(name='Bob', group_id='group', labels=['Entity'])
    indexes = DedupCandidateIndexes(
        existing_nodes=[],
        nodes_by_uuid={},
        normalized_existing=defaultdict(list),
        shingles_by_candidate={},
        lsh_buckets=defaultdict(list),
    )
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[])

    _resolve_with_similarity([extracted], indexes, state)

    assert state.resolved_nodes[0] is None
    assert state.unresolved_indices == [0]
    assert state.duplicate_pairs == []


def test_resolve_with_similarity_multiple_exact_matches_defers_to_llm():
    candidate1 = EntityNode(name='Johnny Appleseed', group_id='group', labels=['Entity'])
    candidate2 = EntityNode(name='Johnny Appleseed', group_id='group', labels=['Entity'])
    extracted = EntityNode(name='Johnny Appleseed', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate1, candidate2])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[])

    _resolve_with_similarity([extracted], indexes, state)

    assert state.resolved_nodes[0] is None
    assert state.unresolved_indices == [0]
    assert state.duplicate_pairs == []


@pytest.mark.asyncio
async def test_resolve_with_llm_candidate_attributes_cannot_overwrite_candidate_id(monkeypatch):
    """Ensure candidate.attributes with a 'candidate_id' key cannot corrupt the LLM context."""
    candidate = EntityNode(name='Dizzy Gillespie', group_id='group', labels=['Entity'])
    candidate.attributes = {'candidate_id': 999, 'genre': 'jazz'}

    extracted = EntityNode(name='Dizzy', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[0])

    captured_context = {}

    def fake_prompt_nodes(context):
        captured_context.update(context)
        return ['prompt']

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations.prompt_library.dedupe_nodes.nodes',
        fake_prompt_nodes,
    )

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={
            'entity_resolutions': [
                {'id': 0, 'name': 'Dizzy Gillespie', 'duplicate_candidate_id': 0}
            ]
        }
    )

    await _resolve_with_llm(
        llm_client,
        [extracted],
        indexes,
        state,
        episode=_make_episode(),
        previous_episodes=[],
        entity_types=None,
    )

    # candidate_id must be the positional index (0), not the adversarial attribute (999)
    assert captured_context['existing_nodes'][0]['candidate_id'] == 0
    # non-colliding attributes should still be present
    assert captured_context['existing_nodes'][0]['genre'] == 'jazz'
    assert state.resolved_nodes[0].uuid == candidate.uuid


@pytest.mark.asyncio
async def test_resolve_with_llm_updates_unresolved(monkeypatch):
    extracted = EntityNode(name='Dizzy', group_id='group', labels=['Entity'])
    candidate = EntityNode(name='Dizzy Gillespie', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[0])

    captured_context = {}

    def fake_prompt_nodes(context):
        captured_context.update(context)
        return ['prompt']

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations.prompt_library.dedupe_nodes.nodes',
        fake_prompt_nodes,
    )

    async def fake_generate_response(*_, **__):
        return {
            'entity_resolutions': [
                {
                    'id': 0,
                    'name': 'Dizzy Gillespie',
                    'duplicate_candidate_id': 0,
                }
            ]
        }

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(side_effect=fake_generate_response)

    await _resolve_with_llm(
        llm_client,
        [extracted],
        indexes,
        state,
        episode=_make_episode(),
        previous_episodes=[],
        entity_types=None,
    )

    assert state.resolved_nodes[0].uuid == candidate.uuid
    assert state.uuid_map[extracted.uuid] == candidate.uuid
    assert isinstance(captured_context['existing_nodes'], list)
    assert captured_context['existing_nodes'][0]['candidate_id'] == 0
    assert (
        captured_context['extracted_nodes'][0]['entity_type_description'] == 'Default Entity Type'
    )
    assert state.duplicate_pairs == [(extracted, candidate)]


@pytest.mark.asyncio
async def test_resolve_with_llm_promotes_generic_candidate_type(monkeypatch):
    extracted = EntityNode(name='Audrey', group_id='group', labels=['Entity', 'Person'])
    candidate = EntityNode(name='Audrey', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[0])

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations.prompt_library.dedupe_nodes.nodes',
        lambda context: ['prompt'],
    )

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={
            'entity_resolutions': [
                {
                    'id': 0,
                    'name': 'Audrey',
                    'duplicate_candidate_id': 0,
                }
            ]
        }
    )

    await _resolve_with_llm(
        llm_client,
        [extracted],
        indexes,
        state,
        episode=_make_episode(),
        previous_episodes=[],
        entity_types=None,
    )

    assert state.resolved_nodes[0].uuid == candidate.uuid
    assert set(state.resolved_nodes[0].labels) == {'Entity', 'Person'}
    assert set(candidate.labels) == {'Entity', 'Person'}
    assert state.uuid_map[extracted.uuid] == candidate.uuid
    assert state.duplicate_pairs == [(extracted, candidate)]


@pytest.mark.asyncio
async def test_resolve_with_llm_ignores_out_of_range_relative_ids(monkeypatch, caplog):
    extracted = EntityNode(name='Dexter', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[0])

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations.prompt_library.dedupe_nodes.nodes',
        lambda context: ['prompt'],
    )

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={
            'entity_resolutions': [
                {
                    'id': 5,
                    'name': 'Dexter',
                    'duplicate_candidate_id': -1,
                }
            ]
        }
    )

    with caplog.at_level(logging.WARNING):
        await _resolve_with_llm(
            llm_client,
            [extracted],
            indexes,
            state,
            episode=_make_episode(),
            previous_episodes=[],
            entity_types=None,
        )

    assert state.resolved_nodes[0] is None
    assert 'Skipping invalid LLM dedupe id 5' in caplog.text


@pytest.mark.asyncio
async def test_resolve_with_llm_ignores_duplicate_relative_ids(monkeypatch):
    extracted = EntityNode(name='Dizzy', group_id='group', labels=['Entity'])
    candidate = EntityNode(name='Dizzy Gillespie', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([candidate])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[0])

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations.prompt_library.dedupe_nodes.nodes',
        lambda context: ['prompt'],
    )

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={
            'entity_resolutions': [
                {
                    'id': 0,
                    'name': 'Dizzy Gillespie',
                    'duplicate_candidate_id': 0,
                },
                {
                    'id': 0,
                    'name': 'Dizzy',
                    'duplicate_candidate_id': -1,
                },
            ]
        }
    )

    await _resolve_with_llm(
        llm_client,
        [extracted],
        indexes,
        state,
        episode=_make_episode(),
        previous_episodes=[],
        entity_types=None,
    )

    assert state.resolved_nodes[0].uuid == candidate.uuid
    assert state.uuid_map[extracted.uuid] == candidate.uuid
    assert state.duplicate_pairs == [(extracted, candidate)]


@pytest.mark.asyncio
async def test_resolve_with_llm_invalid_candidate_id_defaults_to_extracted(monkeypatch):
    extracted = EntityNode(name='Dexter', group_id='group', labels=['Entity'])

    indexes = _build_candidate_indexes([])
    state = DedupResolutionState(resolved_nodes=[None], uuid_map={}, unresolved_indices=[0])

    monkeypatch.setattr(
        'graphiti_core.utils.maintenance.node_operations.prompt_library.dedupe_nodes.nodes',
        lambda context: ['prompt'],
    )

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={
            'entity_resolutions': [
                {
                    'id': 0,
                    'name': 'Dexter',
                    'duplicate_candidate_id': 999,
                }
            ]
        }
    )

    await _resolve_with_llm(
        llm_client,
        [extracted],
        indexes,
        state,
        episode=_make_episode(),
        previous_episodes=[],
        entity_types=None,
    )

    assert state.resolved_nodes[0] == extracted
    assert state.uuid_map[extracted.uuid] == extracted.uuid
    assert state.duplicate_pairs == []


@pytest.mark.asyncio
async def test_batch_summaries_short_summary_no_llm():
    """Test that short summaries are kept as-is without LLM call (optimization)."""
    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={'summaries': [{'name': 'Test Node', 'summary': 'Generated summary'}]}
    )

    node = EntityNode(name='Test Node', group_id='group', labels=['Entity'], summary='Old summary')
    episode = _make_episode()

    await _extract_entity_summaries_batch(
        llm_client,
        [node],
        episode=episode,
        previous_episodes=[],
        should_summarize_node=None,
        edges_by_node={},
    )

    # Short summary should be kept as-is without LLM call
    assert node.summary == 'Old summary'
    # LLM should NOT have been called (summary is short enough)
    llm_client.generate_response.assert_not_awaited()


@pytest.mark.asyncio
async def test_batch_summaries_callback_skip_summary():
    """Test that summary is NOT regenerated when callback returns False."""
    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={'summaries': [{'name': 'Test Node', 'summary': 'This should not be used'}]}
    )

    node = EntityNode(name='Test Node', group_id='group', labels=['Entity'], summary='Old summary')
    episode = _make_episode()

    # Callback that always returns False (skip summary generation)
    async def skip_summary_filter(n: EntityNode) -> bool:
        return False

    await _extract_entity_summaries_batch(
        llm_client,
        [node],
        episode=episode,
        previous_episodes=[],
        should_summarize_node=skip_summary_filter,
        edges_by_node={},
    )

    # Summary should remain unchanged
    assert node.summary == 'Old summary'
    # LLM should NOT have been called for summary
    llm_client.generate_response.assert_not_awaited()


@pytest.mark.asyncio
async def test_batch_summaries_selective_callback():
    """Test callback that selectively skips summaries based on node properties."""
    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(return_value={'summaries': []})

    user_node = EntityNode(name='User', group_id='group', labels=['Entity', 'User'], summary='Old')
    topic_node = EntityNode(
        name='Topic', group_id='group', labels=['Entity', 'Topic'], summary='Old'
    )

    episode = _make_episode()

    # Callback that skips User nodes but generates for others
    async def selective_filter(n: EntityNode) -> bool:
        return 'User' not in n.labels

    await _extract_entity_summaries_batch(
        llm_client,
        [user_node, topic_node],
        episode=episode,
        previous_episodes=[],
        should_summarize_node=selective_filter,
        edges_by_node={},
    )

    # User summary should remain unchanged (callback returned False)
    assert user_node.summary == 'Old'
    # Topic summary should also remain unchanged (short summary optimization)
    assert topic_node.summary == 'Old'
    # LLM should NOT have been called (summaries are short enough)
    llm_client.generate_response.assert_not_awaited()


@pytest.mark.asyncio
async def test_extract_attributes_from_nodes_with_callback():
    """Test that callback is properly passed through extract_attributes_from_nodes."""
    clients, _ = _make_clients()
    clients.llm_client.generate_response = AsyncMock(return_value={'summaries': []})
    clients.embedder.create = AsyncMock(return_value=[0.1, 0.2, 0.3])
    clients.embedder.create_batch = AsyncMock(return_value=[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])

    node1 = EntityNode(name='Node1', group_id='group', labels=['Entity', 'User'], summary='Old1')
    node2 = EntityNode(name='Node2', group_id='group', labels=['Entity', 'Topic'], summary='Old2')

    episode = _make_episode()

    call_tracker = []

    # Callback that tracks which nodes it's called with
    async def tracking_filter(n: EntityNode) -> bool:
        call_tracker.append(n.name)
        return 'User' not in n.labels

    results = await extract_attributes_from_nodes(
        clients,
        [node1, node2],
        episode=episode,
        previous_episodes=[],
        entity_types=None,
        should_summarize_node=tracking_filter,
    )

    # Callback should have been called for both nodes
    assert len(call_tracker) == 2
    assert 'Node1' in call_tracker
    assert 'Node2' in call_tracker

    # Both nodes should keep old summaries (short summary optimization skips LLM)
    node1_result = next(n for n in results if n.name == 'Node1')
    node2_result = next(n for n in results if n.name == 'Node2')

    assert node1_result.summary == 'Old1'
    assert node2_result.summary == 'Old2'


@pytest.mark.asyncio
async def test_batch_summaries_calls_llm_for_long_summary():
    """Test that LLM is called when summary exceeds character limit."""
    from graphiti_core.edges import EntityEdge
    from graphiti_core.utils.text_utils import MAX_SUMMARY_CHARS

    llm_client = MagicMock()
    llm_client.generate_response = AsyncMock(
        return_value={'summaries': [{'name': 'Test Node', 'summary': 'Condensed summary'}]}
    )

    node = EntityNode(name='Test Node', group_id='group', labels=['Entity'], summary='Short')
    episode = _make_episode()

    # Create edges with long facts that exceed the threshold
    long_fact = 'x' * (MAX_SUMMARY_CHARS * 2)
    edge = EntityEdge(
        uuid='edge1',
        group_id='group',
        source_node_uuid=node.uuid,
        target_node_uuid='other-uuid',
        name='test_edge',
        fact=long_fact,
        created_at=utc_now(),
    )

    edges_by_node = {node.uuid: [edge, edge]}  # Multiple long edges

    await _extract_entity_summaries_batch(
        llm_client,
        [node],
        episode=episode,
        previous_episodes=[],
        should_summarize_node=None,
        edges_by_node=edges_by_node,
    )

    # LLM should have been called to condense the long summary
    llm_client.generate_response.assert_awaited_once()
    assert node.summary == 'Condensed summary'