from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest from pydantic import BaseModel from graphiti_core.edges import EntityEdge from graphiti_core.nodes import EntityNode, EpisodicNode from graphiti_core.search.search_config import SearchResults from graphiti_core.utils.maintenance.edge_operations import ( extract_edges, resolve_extracted_edge, resolve_extracted_edges, ) @pytest.fixture def mock_llm_client(): client = MagicMock() client.generate_response = AsyncMock() return client @pytest.fixture def mock_extracted_edge(): return EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='test_edge', group_id='group_1', fact='Test fact', episodes=['episode_1'], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) @pytest.fixture def mock_related_edges(): return [ EntityEdge( source_node_uuid='source_uuid_2', target_node_uuid='target_uuid_2', name='related_edge', group_id='group_1', fact='Related fact', episodes=['episode_2'], created_at=datetime.now(timezone.utc) - timedelta(days=1), valid_at=datetime.now(timezone.utc) - timedelta(days=1), invalid_at=None, ) ] @pytest.fixture def mock_existing_edges(): return [ EntityEdge( source_node_uuid='source_uuid_3', target_node_uuid='target_uuid_3', name='existing_edge', group_id='group_1', fact='Existing fact', episodes=['episode_3'], created_at=datetime.now(timezone.utc) - timedelta(days=2), valid_at=datetime.now(timezone.utc) - timedelta(days=2), invalid_at=None, ) ] @pytest.fixture def mock_current_episode(): return EpisodicNode( uuid='episode_1', content='Current episode content', valid_at=datetime.now(timezone.utc), name='Current Episode', group_id='group_1', source='message', source_description='Test source description', ) @pytest.fixture def mock_previous_episodes(): return [ EpisodicNode( uuid='episode_2', content='Previous episode content', valid_at=datetime.now(timezone.utc) - timedelta(days=1), name='Previous Episode', group_id='group_1', source='message', source_description='Test source description', ) ] # Run the tests if __name__ == '__main__': pytest.main([__file__]) @pytest.mark.asyncio async def test_resolve_extracted_edge_exact_fact_short_circuit( mock_llm_client, mock_existing_edges, mock_current_episode, ): extracted = EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='test_edge', group_id='group_1', fact='Related fact', episodes=['episode_1'], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) related_edges = [ EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='related_edge', group_id='group_1', fact=' related FACT ', episodes=['episode_2'], created_at=datetime.now(timezone.utc) - timedelta(days=1), valid_at=None, invalid_at=None, ) ] resolved_edge, duplicate_edges, invalidated = await resolve_extracted_edge( mock_llm_client, extracted, related_edges, mock_existing_edges, mock_current_episode, edge_type_candidates=None, ) assert resolved_edge is related_edges[0] assert resolved_edge.episodes.count(mock_current_episode.uuid) == 1 assert duplicate_edges == [] assert invalidated == [] mock_llm_client.generate_response.assert_not_called() class OccurredAtEdge(BaseModel): """Edge model stub for OCCURRED_AT.""" @pytest.mark.asyncio async def test_resolve_extracted_edges_keeps_unknown_names(monkeypatch): from graphiti_core.utils.maintenance import edge_operations as edge_ops monkeypatch.setattr(edge_ops, 'create_entity_edge_embeddings', AsyncMock(return_value=None)) monkeypatch.setattr(EntityEdge, 'get_between_nodes', AsyncMock(return_value=[])) async def immediate_gather(*aws, max_coroutines=None): return [await aw for aw in aws] monkeypatch.setattr(edge_ops, 'semaphore_gather', immediate_gather) monkeypatch.setattr(edge_ops, 'search', AsyncMock(return_value=SearchResults())) llm_client = MagicMock() llm_client.generate_response = AsyncMock( return_value={ 'duplicate_facts': [], 'contradicted_facts': [], } ) clients = SimpleNamespace( driver=MagicMock(), llm_client=llm_client, embedder=MagicMock(), cross_encoder=MagicMock(), ) source_node = EntityNode( uuid='source_uuid', name='User Node', group_id='group_1', labels=['User'], ) target_node = EntityNode( uuid='target_uuid', name='Topic Node', group_id='group_1', labels=['Topic'], ) extracted_edge = EntityEdge( source_node_uuid=source_node.uuid, target_node_uuid=target_node.uuid, name='INTERACTED_WITH', group_id='group_1', fact='User interacted with topic', episodes=[], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) episode = EpisodicNode( uuid='episode_uuid', name='Episode', group_id='group_1', source='message', source_description='desc', content='Episode content', valid_at=datetime.now(timezone.utc), ) edge_types = {'OCCURRED_AT': OccurredAtEdge} edge_type_map = {('Event', 'Entity'): ['OCCURRED_AT']} resolved_edges, invalidated_edges, new_edges = await resolve_extracted_edges( clients, [extracted_edge], episode, [source_node, target_node], edge_types, edge_type_map, ) assert resolved_edges[0].name == 'INTERACTED_WITH' assert invalidated_edges == [] assert new_edges == resolved_edges # No duplicates, so all edges are new @pytest.mark.asyncio async def test_resolve_extracted_edge_uses_integer_indices_for_duplicates(mock_llm_client): """Test that resolve_extracted_edge correctly uses integer indices for LLM duplicate detection.""" # Mock LLM to return duplicate_facts with integer indices mock_llm_client.generate_response.return_value = { 'duplicate_facts': [0, 1], # LLM identifies first two related edges as duplicates 'contradicted_facts': [], } extracted_edge = EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='test_edge', group_id='group_1', fact='User likes yoga', episodes=[], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) episode = EpisodicNode( uuid='episode_uuid', name='Episode', group_id='group_1', source='message', source_description='desc', content='Episode content', valid_at=datetime.now(timezone.utc), ) # Create multiple related edges - LLM should receive these with integer indices related_edge_0 = EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='test_edge', group_id='group_1', fact='User enjoys yoga', episodes=['episode_1'], created_at=datetime.now(timezone.utc) - timedelta(days=1), valid_at=None, invalid_at=None, ) related_edge_1 = EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='test_edge', group_id='group_1', fact='User practices yoga', episodes=['episode_2'], created_at=datetime.now(timezone.utc) - timedelta(days=2), valid_at=None, invalid_at=None, ) related_edge_2 = EntityEdge( source_node_uuid='source_uuid', target_node_uuid='target_uuid', name='test_edge', group_id='group_1', fact='User loves swimming', episodes=['episode_3'], created_at=datetime.now(timezone.utc) - timedelta(days=3), valid_at=None, invalid_at=None, ) related_edges = [related_edge_0, related_edge_1, related_edge_2] resolved_edge, invalidated, duplicates = await resolve_extracted_edge( mock_llm_client, extracted_edge, related_edges, [], episode, edge_type_candidates=None, ) # Verify LLM was called mock_llm_client.generate_response.assert_called_once() # Verify the system correctly identified duplicates using integer indices # The LLM returned [0, 1], so related_edge_0 and related_edge_1 should be marked as duplicates assert len(duplicates) == 2 assert related_edge_0 in duplicates assert related_edge_1 in duplicates assert invalidated == [] # Verify that the resolved edge is one of the duplicates (the first one found) # Check UUID since the episode list gets modified assert resolved_edge.uuid == related_edge_0.uuid assert episode.uuid in resolved_edge.episodes @pytest.mark.asyncio async def test_resolve_extracted_edges_fast_path_deduplication(monkeypatch): """Test that resolve_extracted_edges deduplicates exact matches before parallel processing.""" from graphiti_core.utils.maintenance import edge_operations as edge_ops monkeypatch.setattr(edge_ops, 'create_entity_edge_embeddings', AsyncMock(return_value=None)) monkeypatch.setattr(EntityEdge, 'get_between_nodes', AsyncMock(return_value=[])) # Track how many times resolve_extracted_edge is called resolve_call_count = 0 async def mock_resolve_extracted_edge( llm_client, extracted_edge, related_edges, existing_edges, episode, edge_type_candidates=None, ): nonlocal resolve_call_count resolve_call_count += 1 return extracted_edge, [], [] # Mock semaphore_gather to execute awaitable immediately async def immediate_gather(*aws, max_coroutines=None): results = [] for aw in aws: results.append(await aw) return results monkeypatch.setattr(edge_ops, 'semaphore_gather', immediate_gather) monkeypatch.setattr(edge_ops, 'search', AsyncMock(return_value=SearchResults())) monkeypatch.setattr(edge_ops, 'resolve_extracted_edge', mock_resolve_extracted_edge) llm_client = MagicMock() clients = SimpleNamespace( driver=MagicMock(), llm_client=llm_client, embedder=MagicMock(), cross_encoder=MagicMock(), ) source_node = EntityNode( uuid='source_uuid', name='Assistant', group_id='group_1', labels=['Entity'], ) target_node = EntityNode( uuid='target_uuid', name='User', group_id='group_1', labels=['Entity'], ) # Create 3 identical edges edge1 = EntityEdge( source_node_uuid=source_node.uuid, target_node_uuid=target_node.uuid, name='recommends', group_id='group_1', fact='assistant recommends yoga poses', episodes=[], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) edge2 = EntityEdge( source_node_uuid=source_node.uuid, target_node_uuid=target_node.uuid, name='recommends', group_id='group_1', fact=' Assistant Recommends YOGA Poses ', # Different whitespace/case episodes=[], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) edge3 = EntityEdge( source_node_uuid=source_node.uuid, target_node_uuid=target_node.uuid, name='recommends', group_id='group_1', fact='assistant recommends yoga poses', episodes=[], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, ) episode = EpisodicNode( uuid='episode_uuid', name='Episode', group_id='group_1', source='message', source_description='desc', content='Episode content', valid_at=datetime.now(timezone.utc), ) resolved_edges, invalidated_edges, new_edges = await resolve_extracted_edges( clients, [edge1, edge2, edge3], episode, [source_node, target_node], {}, {}, ) # Fast path should have deduplicated the 3 identical edges to 1 # So resolve_extracted_edge should only be called once assert resolve_call_count == 1 assert len(resolved_edges) == 1 assert invalidated_edges == [] assert new_edges == resolved_edges # All edges are new (no graph duplicates) class InterpersonalRelationship(BaseModel): """A relationship between two people.""" class LocatedIn(BaseModel): """A relationship indicating something is located in a place.""" def test_edge_type_signatures_map_preserves_multiple_signatures(): """Test that edge types used across multiple node type pairs preserve all signatures. This tests the fix for the bug where dict comprehension would overwrite previous signatures when the same edge type appeared in multiple node pairs. """ # Edge type map where the same edge type is used for multiple node pair signatures # This is the scenario that was broken before the fix edge_type_map: dict[tuple[str, str], list[str]] = { ('Person', 'Person'): ['InterpersonalRelationship'], ('Person', 'Entity'): ['InterpersonalRelationship'], # Same type, different signature ('Person', 'City'): ['LocatedIn'], ('Entity', 'City'): ['LocatedIn'], # Same type, different signature } edge_types: dict[str, type[BaseModel]] = { 'InterpersonalRelationship': InterpersonalRelationship, 'LocatedIn': LocatedIn, } # Build the mapping the same way as in extract_edges (the fixed implementation) edge_type_signatures_map: dict[str, list[tuple[str, str]]] = {} for signature, edge_type_names in edge_type_map.items(): for edge_type in edge_type_names: if edge_type not in edge_type_signatures_map: edge_type_signatures_map[edge_type] = [] edge_type_signatures_map[edge_type].append(signature) # Verify InterpersonalRelationship has BOTH signatures preserved assert 'InterpersonalRelationship' in edge_type_signatures_map interpersonal_signatures = edge_type_signatures_map['InterpersonalRelationship'] assert len(interpersonal_signatures) == 2 assert ('Person', 'Person') in interpersonal_signatures assert ('Person', 'Entity') in interpersonal_signatures # Verify LocatedIn has BOTH signatures preserved assert 'LocatedIn' in edge_type_signatures_map located_signatures = edge_type_signatures_map['LocatedIn'] assert len(located_signatures) == 2 assert ('Person', 'City') in located_signatures assert ('Entity', 'City') in located_signatures # Verify the edge_types_context structure edge_types_context = [ { 'fact_type_name': type_name, 'fact_type_signatures': edge_type_signatures_map.get(type_name, [('Entity', 'Entity')]), 'fact_type_description': type_model.__doc__, } for type_name, type_model in edge_types.items() ] # Verify the context has the correct structure with plural 'fact_type_signatures' for ctx in edge_types_context: assert 'fact_type_signatures' in ctx assert isinstance(ctx['fact_type_signatures'], list) assert len(ctx['fact_type_signatures']) == 2 # Each type has 2 signatures def test_edge_type_signatures_map_single_signature_still_works(): """Test that edge types with a single signature still work correctly.""" edge_type_map: dict[tuple[str, str], list[str]] = { ('Person', 'Organization'): ['WorksAt'], ('Person', 'City'): ['LivesIn'], } edge_types: dict[str, type[BaseModel]] = { 'WorksAt': BaseModel, 'LivesIn': BaseModel, } # Build the mapping edge_type_signatures_map: dict[str, list[tuple[str, str]]] = {} for signature, edge_type_names in edge_type_map.items(): for edge_type in edge_type_names: if edge_type not in edge_type_signatures_map: edge_type_signatures_map[edge_type] = [] edge_type_signatures_map[edge_type].append(signature) # Verify each edge type has exactly one signature assert len(edge_type_signatures_map['WorksAt']) == 1 assert ('Person', 'Organization') in edge_type_signatures_map['WorksAt'] assert len(edge_type_signatures_map['LivesIn']) == 1 assert ('Person', 'City') in edge_type_signatures_map['LivesIn'] # Verify the context structure edge_types_context = [ { 'fact_type_name': type_name, 'fact_type_signatures': edge_type_signatures_map.get(type_name, [('Entity', 'Entity')]), 'fact_type_description': type_model.__doc__, } for type_name, type_model in edge_types.items() ] for ctx in edge_types_context: assert 'fact_type_signatures' in ctx assert isinstance(ctx['fact_type_signatures'], list) assert len(ctx['fact_type_signatures']) == 1 @pytest.mark.asyncio async def test_extract_edges_drops_self_edges(monkeypatch): """Self-edges (source == target) are dropped during extraction.""" from graphiti_core.prompts.extract_edges import Edge as ExtractedEdge from graphiti_core.prompts.extract_edges import ExtractedEdges alice = EntityNode( uuid='alice_uuid', name='Alice', group_id='group_1', labels=['Person'], ) bob = EntityNode( uuid='bob_uuid', name='Bob', group_id='group_1', labels=['Person'], ) # LLM returns one valid edge and one self-edge llm_response = ExtractedEdges( edges=[ ExtractedEdge( source_entity_name='Alice', target_entity_name='Bob', relation_type='CONGRATULATED', fact='Alice congratulated Bob', valid_at=None, invalid_at=None, ), ExtractedEdge( source_entity_name='Alice', target_entity_name='Alice', relation_type='FEELS_HAPPY', fact='Alice feels happy', valid_at=None, invalid_at=None, ), ] ).model_dump() mock_llm = MagicMock() mock_llm.generate_response = AsyncMock(return_value=llm_response) clients = SimpleNamespace( driver=MagicMock(), llm_client=mock_llm, embedder=MagicMock(), cross_encoder=MagicMock(), ) episode = EpisodicNode( uuid='ep_uuid', name='Episode', group_id='group_1', source='message', source_description='desc', content='Alice congratulated Bob. Alice feels happy.', valid_at=datetime.now(timezone.utc), ) edges = await extract_edges( clients, episode, [alice, bob], [], {}, group_id='group_1', ) # Only the Alice -> Bob edge survives; the Alice -> Alice self-edge is dropped assert len(edges) == 1 assert edges[0].source_node_uuid == 'alice_uuid' assert edges[0].target_node_uuid == 'bob_uuid' assert edges[0].name == 'CONGRATULATED' @pytest.mark.asyncio async def test_extract_edges_keeps_valid_edges_with_same_name_different_nodes(monkeypatch): """Edges between different nodes that happen to share a name are NOT self-edges.""" from graphiti_core.prompts.extract_edges import Edge as ExtractedEdge from graphiti_core.prompts.extract_edges import ExtractedEdges alice = EntityNode( uuid='alice_uuid', name='Alice', group_id='group_1', labels=['Person'], ) paris = EntityNode( uuid='paris_uuid', name='Paris', group_id='group_1', labels=['City'], ) llm_response = ExtractedEdges( edges=[ ExtractedEdge( source_entity_name='Alice', target_entity_name='Paris', relation_type='LIVES_IN', fact='Alice lives in Paris', valid_at=None, invalid_at=None, ), ] ).model_dump() mock_llm = MagicMock() mock_llm.generate_response = AsyncMock(return_value=llm_response) clients = SimpleNamespace( driver=MagicMock(), llm_client=mock_llm, embedder=MagicMock(), cross_encoder=MagicMock(), ) episode = EpisodicNode( uuid='ep_uuid', name='Episode', group_id='group_1', source='message', source_description='desc', content='Alice lives in Paris.', valid_at=datetime.now(timezone.utc), ) edges = await extract_edges( clients, episode, [alice, paris], [], {}, group_id='group_1', ) assert len(edges) == 1 assert edges[0].source_node_uuid == 'alice_uuid' assert edges[0].target_node_uuid == 'paris_uuid' class _EmploymentEdge(BaseModel): """Edge attribute schema mirroring a customer's EMPLOYMENT relation.""" title: str | None = None is_current: str | None = None @pytest.mark.asyncio async def test_resolve_extracted_edge_overcap_attribute_preserves_prior(monkeypatch): """End-to-end: when the LLM bleeds meta-reasoning into one attribute, the cap drops that field and the merge logic falls back to the prior on-edge value — without affecting fields the LLM legitimately updated.""" from graphiti_core.utils.maintenance import edge_operations as edge_ops # Avoid the timestamps LLM call that follows attribute extraction. monkeypatch.setattr(edge_ops, '_extract_edge_timestamps', AsyncMock(return_value=None)) # The LLM returns a clean update for `title` but bleeds 9 KB of meta-reasoning into # `is_current` — the kind of failure the customer reported. bleed = ( 'true (implied by context, but no new information explicitly stated. ' 'I should ideally remove it or set it to null. However, the instruction ' 'is to preserve existing values...) ' ) * 30 llm_client = MagicMock() llm_client.generate_response = AsyncMock( return_value={'title': 'Senior Engineer', 'is_current': bleed} ) extracted_edge = EntityEdge( source_node_uuid='person_uuid', target_node_uuid='company_uuid', name='EMPLOYMENT', group_id='group_1', fact='Sam was promoted to Senior Engineer at Northwind', episodes=[], created_at=datetime.now(timezone.utc), valid_at=None, invalid_at=None, attributes={'title': 'Engineer', 'is_current': 'true'}, ) episode = EpisodicNode( uuid='episode_uuid', name='Episode', group_id='group_1', source='message', source_description='desc', content='Sam was promoted to Senior Engineer.', valid_at=datetime.now(timezone.utc), ) resolved, dupes, invalidated = await resolve_extracted_edge( llm_client, extracted_edge, related_edges=[], existing_edges=[], episode=episode, edge_type_candidates={'EMPLOYMENT': _EmploymentEdge}, ) # Title should reflect the legitimate LLM update. assert resolved.attributes['title'] == 'Senior Engineer' # is_current was bleed-dropped, so the prior value must be preserved (not the # 9 KB rant, not None, not absent). assert resolved.attributes['is_current'] == 'true' assert dupes == [] assert invalidated == []