graphiti-core 0.30.0rc3__tar.gz → 0.30.0rc4__tar.gz

{graphiti_core-0.30.0rc3 → graphiti_core-0.30.0rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.30.0rc3
+Version: 0.30.0rc4
 Summary: A temporal graph building library
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
 Project-URL: Repository, https://github.com/getzep/graphiti

{graphiti_core-0.30.0rc3 → graphiti_core-0.30.0rc4}/graphiti_core/utils/maintenance/edge_operations.py

@@ -43,6 +43,8 @@ from graphiti_core.search.search_filters import SearchFilters
 from graphiti_core.utils.datetime_utils import ensure_utc, utc_now
 from graphiti_core.utils.maintenance.dedup_helpers import _normalize_string_exact
 
+DEFAULT_EDGE_NAME = 'RELATES_TO'
+
 logger = logging.getLogger(__name__)
 
 
@@ -310,6 +312,15 @@ async def resolve_extracted_edges(
 
         edge_types_lst.append(extracted_edge_types)
 
+    for extracted_edge, extracted_edge_types in zip(extracted_edges, edge_types_lst, strict=True):
+        allowed_type_names = set(extracted_edge_types)
+        if not allowed_type_names:
+            if extracted_edge.name != DEFAULT_EDGE_NAME:
+                extracted_edge.name = DEFAULT_EDGE_NAME
+            continue
+        if extracted_edge.name not in allowed_type_names:
+            extracted_edge.name = DEFAULT_EDGE_NAME
+
     # resolve edges with related edges in the graph and find invalidation candidates
     results: list[tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]] = list(
         await semaphore_gather(
@@ -392,7 +403,7 @@ async def resolve_extracted_edge(
     related_edges: list[EntityEdge],
     existing_edges: list[EntityEdge],
     episode: EpisodicNode,
-    edge_types: dict[str, type[BaseModel]] | None = None,
+    edge_type_candidates: dict[str, type[BaseModel]] | None = None,
     ensure_ascii: bool = True,
 ) -> tuple[EntityEdge, list[EntityEdge], list[EntityEdge]]:
     if len(related_edges) == 0 and len(existing_edges) == 0:
@@ -429,9 +440,9 @@ async def resolve_extracted_edge(
                 'fact_type_name': type_name,
                 'fact_type_description': type_model.__doc__,
             }
-            for i, (type_name, type_model) in enumerate(edge_types.items())
+            for i, (type_name, type_model) in enumerate(edge_type_candidates.items())
         ]
-        if edge_types is not None
+        if edge_type_candidates is not None
         else []
     )
 
@@ -468,7 +479,8 @@ async def resolve_extracted_edge(
     ]
 
     fact_type: str = response_object.fact_type
-    if fact_type.upper() != 'DEFAULT' and edge_types is not None:
+    candidate_type_names = set(edge_type_candidates or {})
+    if candidate_type_names and fact_type in candidate_type_names:
         resolved_edge.name = fact_type
 
         edge_attributes_context = {
@@ -478,7 +490,7 @@ async def resolve_extracted_edge(
             'ensure_ascii': ensure_ascii,
         }
 
-        edge_model = edge_types.get(fact_type)
+        edge_model = edge_type_candidates.get(fact_type) if edge_type_candidates else None
         if edge_model is not None and len(edge_model.model_fields) != 0:
             edge_attributes_response = await llm_client.generate_response(
                 prompt_library.extract_edges.extract_attributes(edge_attributes_context),
@@ -487,6 +499,9 @@ async def resolve_extracted_edge(
             )
 
             resolved_edge.attributes = edge_attributes_response
+    elif fact_type.upper() != 'DEFAULT':
+        resolved_edge.name = DEFAULT_EDGE_NAME
+        resolved_edge.attributes = {}
 
     end = time()
     logger.debug(

{graphiti_core-0.30.0rc3 → graphiti_core-0.30.0rc4}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.30.0pre3"
+version = "0.30.0pre4"
 authors = [
     { name = "Paul Paliychuk", email = "paul@getzep.com" },
     { name = "Preston Rasmussen", email = "preston@getzep.com" },

graphiti_core-0.30.0rc4/tests/utils/maintenance/test_edge_operations.py

@@ -0,0 +1,291 @@
+from datetime import datetime, timedelta, timezone
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from pydantic import BaseModel
+
+from graphiti_core.edges import EntityEdge
+from graphiti_core.graphiti_types import GraphitiClients
+from graphiti_core.nodes import EntityNode, EpisodicNode
+from graphiti_core.search.search_config import SearchResults
+from graphiti_core.utils.maintenance.edge_operations import (
+    DEFAULT_EDGE_NAME,
+    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,
+        ensure_ascii=True,
+    )
+
+    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_resets_unmapped_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': [], 'fact_type': 'DEFAULT'}
+    )
+
+    clients = GraphitiClients.model_construct(
+        driver=MagicMock(),
+        llm_client=llm_client,
+        embedder=MagicMock(),
+        cross_encoder=MagicMock(),
+        ensure_ascii=True,
+    )
+
+    source_node = EntityNode(
+        uuid='source_uuid',
+        name='Document Node',
+        group_id='group_1',
+        labels=['Document'],
+    )
+    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='OCCURRED_AT',
+        group_id='group_1',
+        fact='Document occurred at somewhere',
+        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 = await resolve_extracted_edges(
+        clients,
+        [extracted_edge],
+        episode,
+        [source_node, target_node],
+        edge_types,
+        edge_type_map,
+    )
+
+    assert resolved_edges[0].name == DEFAULT_EDGE_NAME
+    assert invalidated_edges == []
+
+
+@pytest.mark.asyncio
+async def test_resolve_extracted_edge_rejects_unmapped_fact_type(mock_llm_client):
+    mock_llm_client.generate_response.return_value = {
+        'duplicate_facts': [],
+        'contradicted_facts': [],
+        'fact_type': 'OCCURRED_AT',
+    }
+
+    extracted_edge = EntityEdge(
+        source_node_uuid='source_uuid',
+        target_node_uuid='target_uuid',
+        name='OCCURRED_AT',
+        group_id='group_1',
+        fact='Document occurred at somewhere',
+        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),
+    )
+
+    related_edge = EntityEdge(
+        source_node_uuid='alt_source',
+        target_node_uuid='alt_target',
+        name='OTHER',
+        group_id='group_1',
+        fact='Different fact',
+        episodes=[],
+        created_at=datetime.now(timezone.utc),
+        valid_at=None,
+        invalid_at=None,
+    )
+
+    resolved_edge, duplicates, invalidated = await resolve_extracted_edge(
+        mock_llm_client,
+        extracted_edge,
+        [related_edge],
+        [],
+        episode,
+        edge_type_candidates={},
+        ensure_ascii=True,
+    )
+
+    assert resolved_edge.name == DEFAULT_EDGE_NAME
+    assert duplicates == []
+    assert invalidated == []

{graphiti_core-0.30.0rc3 → graphiti_core-0.30.0rc4}/uv.lock

@@ -783,7 +783,7 @@ wheels = [
 
 [[package]]
 name = "graphiti-core"
-version = "0.30.0rc3"
+version = "0.30.0rc4"
 source = { editable = "." }
 dependencies = [
     { name = "diskcache" },

graphiti_core-0.30.0rc3/tests/utils/maintenance/test_edge_operations.py

@@ -1,144 +0,0 @@
-from datetime import datetime, timedelta, timezone
-from unittest.mock import MagicMock
-
-import pytest
-
-from graphiti_core.edges import EntityEdge
-from graphiti_core.nodes import EpisodicNode
-from graphiti_core.utils.maintenance.edge_operations import resolve_extracted_edge
-
-
-@pytest.fixture
-def mock_llm_client():
-    return MagicMock()
-
-
-@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_types=None,
-        ensure_ascii=True,
-    )
-
-    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()