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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.30.0rc2
+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.0rc2 → graphiti_core-0.30.0rc4}/graphiti_core/prompts/dedupe_nodes.py

@@ -92,12 +92,23 @@ def node(context: dict[str, Any]) -> list[Message]:
 
          TASK:
          1. Compare `new_entity` against each item in `existing_entities`.
-         2. If it refers to the same real‐world object or concept, collect its index.
-         3. Let `duplicate_idx` = the *first* collected index, or –1 if none.
-         4. Let `duplicates` = the list of *all* collected indices (empty list if none).
-        
-        Also return the full name of the NEW ENTITY (whether it is the name of the NEW ENTITY, a node it
-        is a duplicate of, or a combination of the two).
+         2. If it refers to the same real-world object or concept, collect its index.
+         3. Let `duplicate_idx` = the smallest collected index, or -1 if none.
+         4. Let `duplicates` = the sorted list of all collected indices (empty list if none).
+
+        Respond with a JSON object containing an "entity_resolutions" array with a single entry:
+        {{
+            "entity_resolutions": [
+                {{
+                    "id": integer id from NEW ENTITY,
+                    "name": the best full name for the entity,
+                    "duplicate_idx": integer index of the best duplicate in EXISTING ENTITIES, or -1 if none,
+                    "duplicates": sorted list of all duplicate indices you collected (deduplicate the list, use [] when none)
+                }}
+            ]
+        }}
+
+        Only reference indices that appear in EXISTING ENTITIES, and return [] / -1 when unsure.
         """,
         ),
     ]
@@ -126,26 +137,26 @@ def nodes(context: dict[str, Any]) -> list[Message]:
         {{
             id: integer id of the entity,
             name: "name of the entity",
-            entity_type: "ontological classification of the entity",
-            entity_type_description: "Description of what the entity type represents",
-            duplication_candidates: [
-                {{
-                    idx: integer index of the candidate entity,
-                    name: "name of the candidate entity",
-                    entity_type: "ontological classification of the candidate entity",
-                    ...<additional attributes>
-                }}
-            ]
+            entity_type: ["Entity", "<optional additional label>", ...],
+            entity_type_description: "Description of what the entity type represents"
         }}
-        
+
         <ENTITIES>
         {to_prompt_json(context['extracted_nodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
         </ENTITIES>
-        
+
         <EXISTING ENTITIES>
         {to_prompt_json(context['existing_nodes'], ensure_ascii=context.get('ensure_ascii', True), indent=2)}
         </EXISTING ENTITIES>
 
+        Each entry in EXISTING ENTITIES is an object with the following structure:
+        {{
+            idx: integer index of the candidate entity (use this when referencing a duplicate),
+            name: "name of the candidate entity",
+            entity_types: ["Entity", "<optional additional label>", ...],
+            ...<additional attributes such as summaries or metadata>
+        }}
+
         For each of the above ENTITIES, determine if the entity is a duplicate of any of the EXISTING ENTITIES.
 
         Entities should only be considered duplicates if they refer to the *same real-world object or concept*.
@@ -155,14 +166,19 @@ def nodes(context: dict[str, Any]) -> list[Message]:
         - They have similar names or purposes but refer to separate instances or concepts.
 
         Task:
-        Your response will be a list called entity_resolutions which contains one entry for each entity.
-        
-        For each entity, return the id of the entity as id, the name of the entity as name, and the duplicate_idx
-        as an integer.
-        
-        - If an entity is a duplicate of one of the EXISTING ENTITIES, return the idx of the candidate it is a 
-        duplicate of.
-        - If an entity is not a duplicate of one of the EXISTING ENTITIES, return the -1 as the duplication_idx
+        Respond with a JSON object that contains an "entity_resolutions" array with one entry for each entity in ENTITIES, ordered by the entity id.
+
+        For every entity, return an object with the following keys:
+        {{
+            "id": integer id from ENTITIES,
+            "name": the best full name for the entity (preserve the original name unless a duplicate has a more complete name),
+            "duplicate_idx": the idx of the EXISTING ENTITY that is the best duplicate match, or -1 if there is no duplicate,
+            "duplicates": a sorted list of all idx values from EXISTING ENTITIES that refer to duplicates (deduplicate the list, use [] when none or unsure)
+        }}
+
+        - Only use idx values that appear in EXISTING ENTITIES.
+        - Set duplicate_idx to the smallest idx you collected for that entity, or -1 if duplicates is empty.
+        - Never fabricate entities or indices.
         """,
         ),
     ]

{graphiti_core-0.30.0rc2 → 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.0rc2 → graphiti_core-0.30.0rc4}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.30.0pre2"
+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.0rc2 → graphiti_core-0.30.0rc4}/uv.lock

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