graphiti-core 0.20.4__py3-none-any.whl → 0.21.0__py3-none-any.whl

graphiti_core/utils/maintenance/node_operations.py

@@ -15,6 +15,7 @@ limitations under the License.
 """
 
 import logging
+from collections.abc import Awaitable, Callable
 from time import time
 from typing import Any
 
@@ -24,7 +25,12 @@ from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, 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 NodeDuplicate, NodeResolutions
 from graphiti_core.prompts.extract_nodes import (
@@ -38,28 +44,37 @@ 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.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,
+)
 
 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],
-    ensure_ascii: bool = False,
+    group_id: str | None = None,
 ) -> list[str]:
     # Prepare context for LLM
     context = {
         'episode_content': episode.content,
         'previous_episodes': [ep.content for ep in previous_episodes],
         'extracted_entities': node_names,
-        'ensure_ascii': ensure_ascii,
     }
 
     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
     )
     missed_entities = llm_response.get('missed_entities', [])
 
@@ -108,7 +123,6 @@ async def extract_nodes(
         'custom_prompt': custom_prompt,
         'entity_types': entity_types_context,
         'source_description': episode.source_description,
-        'ensure_ascii': clients.ensure_ascii,
     }
 
     while entities_missed and reflexion_iterations <= MAX_REFLEXION_ITERATIONS:
@@ -116,14 +130,19 @@ async def extract_nodes(
             llm_response = await llm_client.generate_response(
                 prompt_library.extract_nodes.extract_message(context),
                 response_model=ExtractedEntities,
+                group_id=episode.group_id,
             )
         elif episode.source == EpisodeType.text:
             llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_text(context), response_model=ExtractedEntities
+                prompt_library.extract_nodes.extract_text(context),
+                response_model=ExtractedEntities,
+                group_id=episode.group_id,
             )
         elif episode.source == EpisodeType.json:
             llm_response = await llm_client.generate_response(
-                prompt_library.extract_nodes.extract_json(context), response_model=ExtractedEntities
+                prompt_library.extract_nodes.extract_json(context),
+                response_model=ExtractedEntities,
+                group_id=episode.group_id,
             )
 
         response_object = ExtractedEntities(**llm_response)
@@ -137,7 +156,7 @@ async def extract_nodes(
                 episode,
                 previous_episodes,
                 [entity.name for entity in extracted_entities],
-                clients.ensure_ascii,
+                episode.group_id,
             )
 
             entities_missed = len(missing_entities) != 0
@@ -178,20 +197,16 @@ async def extract_nodes(
         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, type[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(
@@ -205,33 +220,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]
+
+    if existing_nodes_override is not None:
+        candidate_nodes.extend(existing_nodes_override)
 
-    existing_nodes_dict: dict[str, EntityNode] = {node.uuid: node for node in candidate_nodes}
+    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: list[EntityNode] = list(existing_nodes_dict.values())
+    return ordered_candidates
 
-    existing_nodes_context = (
-        [
-            {
-                **{
-                    'idx': i,
-                    'name': candidate.name,
-                    'entity_types': candidate.labels,
-                },
-                **candidate.attributes,
-            }
-            for i, candidate in enumerate(existing_nodes)
-        ],
-    )
+
+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 {}
 
-    # Prepare context for LLM
+    llm_extracted_nodes = [extracted_nodes[i] for i in state.unresolved_indices]
+
     extracted_nodes_context = [
         {
             'id': i,
@@ -242,17 +267,49 @@ 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 [],
-        'ensure_ascii': clients.ensure_ascii,
+        'previous_episodes': (
+            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
+        ),
     }
 
     llm_response = await llm_client.generate_response(
@@ -262,41 +319,127 @@ async def resolve_extracted_nodes(
 
     node_resolutions: list[NodeDuplicate] = NodeResolutions(**llm_response).entity_resolutions
 
-    resolved_nodes: list[EntityNode] = []
-    uuid_map: dict[str, str] = {}
-    node_duplicates: list[tuple[EntityNode, EntityNode]] = []
+    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
+
+    logger.debug(
+        'Received %d resolutions for %d entities',
+        len(node_resolutions),
+        len(state.unresolved_indices),
+    )
+
+    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),
+        )
+
     for resolution in node_resolutions:
-        resolution_id: int = resolution.id
+        relative_id: int = resolution.id
         duplicate_idx: int = resolution.duplicate_idx
 
-        extracted_node = extracted_nodes[resolution_id]
+        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
 
-        resolved_node = (
-            existing_nodes[duplicate_idx]
-            if 0 <= duplicate_idx < len(existing_nodes)
-            else extracted_node
-        )
+        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_node.name = resolution.get('name')
+    indexes: DedupCandidateIndexes = _build_candidate_indexes(existing_nodes)
 
-        resolved_nodes.append(resolved_node)
-        uuid_map[extracted_node.uuid] = resolved_node.uuid
+    state = DedupResolutionState(
+        resolved_nodes=[None] * len(extracted_nodes),
+        uuid_map={},
+        unresolved_indices=[],
+    )
 
-        duplicates: list[int] = resolution.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
+    _resolve_with_similarity(extracted_nodes, indexes, state)
 
-            node_duplicates.append((extracted_node, existing_node))
+    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(f'Resolved nodes: {[(n.name, n.uuid) for n in resolved_nodes]}')
+    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(
@@ -305,6 +448,7 @@ async def extract_attributes_from_nodes(
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | 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
@@ -315,10 +459,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,
-                clients.ensure_ascii,
+                (
+                    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
         ]
@@ -335,7 +481,7 @@ async def extract_attributes_from_node(
     episode: EpisodicNode | None = None,
     previous_episodes: list[EpisodicNode] | None = None,
     entity_type: type[BaseModel] | None = None,
-    ensure_ascii: bool = False,
+    should_summarize_node: NodeSummaryFilter | None = None,
 ) -> EntityNode:
     node_context: dict[str, Any] = {
         'name': node.name,
@@ -347,43 +493,53 @@ async def extract_attributes_from_node(
     attributes_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 [],
-        'ensure_ascii': ensure_ascii,
+        'previous_episodes': (
+            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
+        ),
     }
 
     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 [],
-        'ensure_ascii': ensure_ascii,
+        'previous_episodes': (
+            [ep.content for ep in previous_episodes] if previous_episodes is not None else []
+        ),
     }
 
+    has_entity_attributes: bool = bool(
+        entity_type is not None and len(entity_type.model_fields) != 0
+    )
+
     llm_response = (
         (
             await llm_client.generate_response(
                 prompt_library.extract_nodes.extract_attributes(attributes_context),
                 response_model=entity_type,
                 model_size=ModelSize.small,
+                group_id=node.group_id,
             )
         )
-        if entity_type is not None
+        if has_entity_attributes
         else {}
     )
 
-    summary_response = await llm_client.generate_response(
-        prompt_library.extract_nodes.extract_summary(summary_context),
-        response_model=EntitySummary,
-        model_size=ModelSize.small,
-    )
+    # Determine if summary should be generated
+    generate_summary = True
+    if should_summarize_node is not None:
+        generate_summary = await should_summarize_node(node)
+
+    # Conditionally generate summary
+    if generate_summary:
+        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,
+        )
+        node.summary = summary_response.get('summary', '')
 
-    if entity_type is not None:
+    if has_entity_attributes and entity_type is not None:
         entity_type(**llm_response)
-
-    node.summary = summary_response.get('summary', '')
     node_attributes = {key: value for key, value in llm_response.items()}
 
     node.attributes.update(node_attributes)

graphiti_core/utils/maintenance/temporal_operations.py

@@ -35,14 +35,12 @@ async def extract_edge_dates(
     edge: EntityEdge,
     current_episode: EpisodicNode,
     previous_episodes: list[EpisodicNode],
-    ensure_ascii: bool = False,
 ) -> tuple[datetime | None, datetime | None]:
     context = {
         'edge_fact': edge.fact,
         'current_episode': current_episode.content,
         'previous_episodes': [ep.content for ep in previous_episodes],
         'reference_timestamp': current_episode.valid_at.isoformat(),
-        'ensure_ascii': ensure_ascii,
     }
     llm_response = await llm_client.generate_response(
         prompt_library.extract_edge_dates.v1(context), response_model=EdgeDates
@@ -75,7 +73,6 @@ async def get_edge_contradictions(
     llm_client: LLMClient,
     new_edge: EntityEdge,
     existing_edges: list[EntityEdge],
-    ensure_ascii: bool = False,
 ) -> list[EntityEdge]:
     start = time()
 
@@ -87,7 +84,6 @@ async def get_edge_contradictions(
     context = {
         'new_edge': new_edge_context,
         'existing_edges': existing_edge_context,
-        'ensure_ascii': ensure_ascii,
     }
 
     llm_response = await llm_client.generate_response(

{graphiti_core-0.20.4.dist-info → graphiti_core-0.21.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: graphiti-core
-Version: 0.20.4
+Version: 0.21.0
 Summary: A temporal graph building library
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
 Project-URL: Repository, https://github.com/getzep/graphiti
@@ -20,6 +20,7 @@ Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
 Provides-Extra: dev
 Requires-Dist: anthropic>=0.49.0; extra == 'dev'
+Requires-Dist: boto3>=1.39.16; extra == 'dev'
 Requires-Dist: diskcache-stubs>=5.6.3.6.20240818; extra == 'dev'
 Requires-Dist: falkordb<2.0.0,>=1.1.2; extra == 'dev'
 Requires-Dist: google-genai>=1.8.0; extra == 'dev'
@@ -28,9 +29,11 @@ Requires-Dist: ipykernel>=6.29.5; extra == 'dev'
 Requires-Dist: jupyterlab>=4.2.4; extra == 'dev'
 Requires-Dist: kuzu>=0.11.2; extra == 'dev'
 Requires-Dist: langchain-anthropic>=0.2.4; extra == 'dev'
+Requires-Dist: langchain-aws>=0.2.29; extra == 'dev'
 Requires-Dist: langchain-openai>=0.2.6; extra == 'dev'
 Requires-Dist: langgraph>=0.2.15; extra == 'dev'
 Requires-Dist: langsmith>=0.1.108; extra == 'dev'
+Requires-Dist: opensearch-py>=3.0.0; extra == 'dev'
 Requires-Dist: pyright>=1.1.404; extra == 'dev'
 Requires-Dist: pytest-asyncio>=0.24.0; extra == 'dev'
 Requires-Dist: pytest-xdist>=3.6.1; extra == 'dev'
@@ -47,6 +50,9 @@ Provides-Extra: groq
 Requires-Dist: groq>=0.2.0; extra == 'groq'
 Provides-Extra: kuzu
 Requires-Dist: kuzu>=0.11.2; extra == 'kuzu'
+Provides-Extra: neo4j-opensearch
+Requires-Dist: boto3>=1.39.16; extra == 'neo4j-opensearch'
+Requires-Dist: opensearch-py>=3.0.0; extra == 'neo4j-opensearch'
 Provides-Extra: neptune
 Requires-Dist: boto3>=1.39.16; extra == 'neptune'
 Requires-Dist: langchain-aws>=0.2.29; extra == 'neptune'