wehub-resource-sync
4a19d70af1
Lint with Ruff / ruff (push) Has been cancelled
MCP Server Tests / live-mcp-tests (push) Has been cancelled
Tests / unit-tests (push) Has been cancelled
Tests / database-integration-tests (push) Has been cancelled
CodeQL Advanced / Analyze (actions) (push) Has been cancelled
CodeQL Advanced / Analyze (python) (push) Has been cancelled
Server Tests / live-server-tests (push) Has been cancelled
Pyright Type Check / pyright (push) Has been cancelled
1033 lines
37 KiB
Python
1033 lines
37 KiB
Python
"""
|
|
Copyright 2024, Zep Software, Inc.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
import logging
|
|
from collections.abc import Awaitable, Callable
|
|
from time import time
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from graphiti_core.edges import EntityEdge
|
|
from graphiti_core.graphiti_types import GraphitiClients
|
|
from graphiti_core.helpers import 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.prompts import prompt_library
|
|
from graphiti_core.prompts.dedupe_nodes import NodeDuplicate, NodeResolutions
|
|
from graphiti_core.prompts.extract_nodes import (
|
|
ExtractedEntities,
|
|
ExtractedEntity,
|
|
SummarizedEntities,
|
|
)
|
|
from graphiti_core.search.search_filters import SearchFilters
|
|
from graphiti_core.search.search_utils import node_similarity_search
|
|
from graphiti_core.utils.datetime_utils import utc_now
|
|
from graphiti_core.utils.maintenance.attribute_utils import apply_capped_attributes
|
|
from graphiti_core.utils.maintenance.dedup_helpers import (
|
|
DedupCandidateIndexes,
|
|
DedupResolutionState,
|
|
_build_candidate_indexes,
|
|
_normalize_string_exact,
|
|
_promote_resolved_node,
|
|
_resolve_with_similarity,
|
|
)
|
|
from graphiti_core.utils.text_utils import (
|
|
MAX_SUMMARY_CHARS,
|
|
concatenate_episodes,
|
|
truncate_at_sentence,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Maximum number of nodes to summarize in a single LLM call
|
|
MAX_NODES = 30
|
|
NODE_DEDUP_CANDIDATE_LIMIT = 15
|
|
NODE_DEDUP_COSINE_MIN_SCORE = 0.6
|
|
|
|
NodeSummaryFilter = Callable[[EntityNode], Awaitable[bool]]
|
|
|
|
|
|
async def extract_nodes(
|
|
clients: GraphitiClients,
|
|
episode: EpisodicNode | list[EpisodicNode],
|
|
previous_episodes: list[EpisodicNode],
|
|
entity_types: dict[str, type[BaseModel]] | None = None,
|
|
excluded_entity_types: list[str] | None = None,
|
|
custom_extraction_instructions: str | None = None,
|
|
) -> tuple[list[EntityNode], dict[str, list[int]]]:
|
|
"""Extract entity nodes from one or more episodes.
|
|
|
|
Parameters
|
|
----------
|
|
episode : EpisodicNode | list[EpisodicNode]
|
|
A single episode or a list of episodes to extract entities from.
|
|
When a list is provided, their contents are concatenated for extraction
|
|
and the first episode is used for metadata (source type, group_id, etc.).
|
|
|
|
Returns
|
|
-------
|
|
tuple[list[EntityNode], dict[str, list[int]]]
|
|
A tuple of (extracted_nodes, node_episode_index_map) where
|
|
node_episode_index_map maps node UUID to a list of 0-indexed episode
|
|
positions that the node was extracted from.
|
|
"""
|
|
episodes = episode if isinstance(episode, list) else [episode]
|
|
primary_episode = episodes[0]
|
|
|
|
start = time()
|
|
llm_client = clients.llm_client
|
|
|
|
# Build entity types context
|
|
entity_types_context = _build_entity_types_context(entity_types)
|
|
|
|
# Build episode attribution instructions for multi-episode extraction
|
|
episode_attribution = ''
|
|
if len(episodes) > 1:
|
|
episode_attribution = (
|
|
'\n7. **Episode Attribution**: The content contains multiple episodes labeled '
|
|
'[Episode 0], [Episode 1], etc. Each episode header includes a timestamp indicating '
|
|
'when that episode occurred. For each extracted entity, set `episode_indices` '
|
|
'to the 0-based list of episode numbers where that entity is mentioned. '
|
|
'An entity appearing in Episodes 0 and 2 should have `episode_indices: [0, 2]`.'
|
|
)
|
|
|
|
# Build base context
|
|
context = {
|
|
'episode_content': concatenate_episodes(episodes),
|
|
'episode_timestamp': primary_episode.valid_at.isoformat(),
|
|
'previous_episodes': [
|
|
{
|
|
'content': ep.content,
|
|
'timestamp': ep.valid_at.isoformat() if ep.valid_at else None,
|
|
}
|
|
for ep in previous_episodes
|
|
],
|
|
'custom_extraction_instructions': (custom_extraction_instructions or '')
|
|
+ episode_attribution,
|
|
'entity_types': entity_types_context,
|
|
'source_description': primary_episode.source_description,
|
|
}
|
|
|
|
# Extract entities
|
|
extracted_entities = await _extract_nodes_single(llm_client, primary_episode, context)
|
|
|
|
# Filter empty names
|
|
filtered_entities = [e for e in extracted_entities if e.name.strip()]
|
|
|
|
end = time()
|
|
logger.debug(f'Extracted {len(filtered_entities)} entities in {(end - start) * 1000:.0f} ms')
|
|
|
|
# Convert to EntityNode objects with episode attribution
|
|
extracted_nodes, node_episode_index_map = _create_entity_nodes(
|
|
filtered_entities, entity_types_context, excluded_entity_types, episodes
|
|
)
|
|
extracted_nodes = _collapse_exact_duplicate_extracted_nodes(
|
|
extracted_nodes, node_episode_index_map
|
|
)
|
|
|
|
logger.debug(f'Extracted nodes: {[n.uuid for n in extracted_nodes]}')
|
|
return extracted_nodes, node_episode_index_map
|
|
|
|
|
|
def _build_entity_types_context(
|
|
entity_types: dict[str, type[BaseModel]] | None,
|
|
) -> list[dict]:
|
|
"""Build entity types context with ID mappings."""
|
|
entity_types_context = [
|
|
{
|
|
'entity_type_id': 0,
|
|
'entity_type_name': 'Entity',
|
|
'entity_type_description': (
|
|
'A specific, identifiable entity that does not fit any of the other listed '
|
|
'types. Must still be a concrete, meaningful thing — specific enough to be '
|
|
'uniquely identifiable. GOOD: a named entity not covered by the other types. '
|
|
'BAD: "luck", "ideas", "tomorrow", "things", "them", "everybody", '
|
|
'"a sense of wonder", "great times". '
|
|
'When in doubt, do not extract the entity.'
|
|
),
|
|
}
|
|
]
|
|
|
|
if entity_types is not None:
|
|
entity_types_context += [
|
|
{
|
|
'entity_type_id': i + 1,
|
|
'entity_type_name': type_name,
|
|
'entity_type_description': type_model.__doc__,
|
|
}
|
|
for i, (type_name, type_model) in enumerate(entity_types.items())
|
|
]
|
|
|
|
return entity_types_context
|
|
|
|
|
|
def _get_entity_type_description(
|
|
labels: list[str], entity_types: dict[str, type[BaseModel]] | None
|
|
) -> str:
|
|
type_name = next((item for item in labels if item != 'Entity'), '')
|
|
type_model = entity_types.get(type_name) if entity_types is not None else None
|
|
return (type_model.__doc__ if type_model is not None else None) or 'Default Entity Type'
|
|
|
|
|
|
def _truncate_type_description(docstring: str) -> str:
|
|
"""Extract a concise type description from a docstring for summary prompts.
|
|
|
|
Returns the first paragraph (up to the first blank line), capped at 3
|
|
sentences. This strips GOOD/BAD examples, trigger patterns, and other
|
|
extraction-specific guidance that is irrelevant to summarization.
|
|
"""
|
|
# Take only the first paragraph.
|
|
paragraph_lines: list[str] = []
|
|
for line in docstring.splitlines():
|
|
if not line.strip():
|
|
if paragraph_lines:
|
|
break
|
|
continue # skip leading blank lines
|
|
paragraph_lines.append(line)
|
|
|
|
text = ' '.join(line.strip() for line in paragraph_lines)
|
|
|
|
# Cap at 3 sentences.
|
|
sentences: list[str] = []
|
|
remaining = text
|
|
for _ in range(3):
|
|
idx = _find_sentence_end(remaining)
|
|
if idx == -1:
|
|
sentences.append(remaining)
|
|
remaining = ''
|
|
break
|
|
sentences.append(remaining[: idx + 1])
|
|
remaining = remaining[idx + 1 :].lstrip()
|
|
return ' '.join(sentences).strip()
|
|
|
|
|
|
def _find_sentence_end(text: str) -> int:
|
|
"""Return the index of the first sentence boundary.
|
|
|
|
A sentence ends at `.`, `!`, or `?` when followed by end-of-string or
|
|
a space then an uppercase letter. This avoids splitting on abbreviations
|
|
like "e.g.", "Dr.", or decimals like "2.0".
|
|
"""
|
|
n = len(text)
|
|
for i, ch in enumerate(text):
|
|
if ch not in '.!?':
|
|
continue
|
|
# End of string counts as a sentence boundary.
|
|
if i + 1 >= n:
|
|
return i
|
|
# Space followed by an uppercase letter is a sentence boundary.
|
|
if text[i + 1] == ' ' and i + 2 < n and text[i + 2].isupper():
|
|
return i
|
|
return -1
|
|
|
|
|
|
async def _extract_nodes_single(
|
|
llm_client: LLMClient,
|
|
episode: EpisodicNode,
|
|
context: dict,
|
|
) -> list[ExtractedEntity]:
|
|
"""Extract entities using a single LLM call."""
|
|
llm_response = await _call_extraction_llm(llm_client, episode, context)
|
|
response_object = ExtractedEntities(**llm_response)
|
|
return response_object.extracted_entities
|
|
|
|
|
|
async def _call_extraction_llm(
|
|
llm_client: LLMClient,
|
|
episode: EpisodicNode,
|
|
context: dict,
|
|
) -> dict:
|
|
"""Call the appropriate extraction prompt based on episode type."""
|
|
if episode.source == EpisodeType.message:
|
|
prompt = prompt_library.extract_nodes.extract_message(context)
|
|
prompt_name = 'extract_nodes.extract_message'
|
|
elif episode.source == EpisodeType.text:
|
|
prompt = prompt_library.extract_nodes.extract_text(context)
|
|
prompt_name = 'extract_nodes.extract_text'
|
|
elif episode.source == EpisodeType.json:
|
|
prompt = prompt_library.extract_nodes.extract_json(context)
|
|
prompt_name = 'extract_nodes.extract_json'
|
|
else:
|
|
# Fallback to text extraction
|
|
prompt = prompt_library.extract_nodes.extract_text(context)
|
|
prompt_name = 'extract_nodes.extract_text'
|
|
|
|
return await llm_client.generate_response(
|
|
prompt,
|
|
response_model=ExtractedEntities,
|
|
group_id=episode.group_id,
|
|
prompt_name=prompt_name,
|
|
)
|
|
|
|
|
|
def _create_entity_nodes(
|
|
extracted_entities: list[ExtractedEntity],
|
|
entity_types_context: list[dict],
|
|
excluded_entity_types: list[str] | None,
|
|
episodes: list[EpisodicNode],
|
|
) -> tuple[list[EntityNode], dict[str, list[int]]]:
|
|
"""Convert ExtractedEntity objects to EntityNode objects.
|
|
|
|
Returns
|
|
-------
|
|
tuple[list[EntityNode], dict[str, list[int]]]
|
|
A tuple of (nodes, node_episode_index_map) where node_episode_index_map
|
|
maps each node UUID to 0-indexed episode positions the node was extracted from.
|
|
"""
|
|
primary_episode = episodes[0]
|
|
extracted_nodes = []
|
|
node_episode_index_map: dict[str, list[int]] = {}
|
|
|
|
for extracted_entity in extracted_entities:
|
|
type_id = extracted_entity.entity_type_id
|
|
if 0 <= type_id < len(entity_types_context):
|
|
entity_type_name = entity_types_context[type_id].get('entity_type_name')
|
|
else:
|
|
entity_type_name = 'Entity'
|
|
|
|
# Check if this entity type should be excluded
|
|
if excluded_entity_types and entity_type_name in excluded_entity_types:
|
|
logger.debug(f'Excluding entity of type "{entity_type_name}"')
|
|
continue
|
|
|
|
labels: list[str] = list({'Entity', str(entity_type_name)})
|
|
|
|
new_node = EntityNode(
|
|
name=extracted_entity.name,
|
|
group_id=primary_episode.group_id,
|
|
labels=labels,
|
|
summary='',
|
|
created_at=utc_now(),
|
|
)
|
|
extracted_nodes.append(new_node)
|
|
|
|
# Map node to 0-indexed episode positions (LLM returns 0-indexed).
|
|
# Clamp to valid range; fall back to all episodes if empty.
|
|
indices = [i for i in extracted_entity.episode_indices if 0 <= i < len(episodes)]
|
|
if not indices:
|
|
indices = list(range(len(episodes)))
|
|
node_episode_index_map[new_node.uuid] = indices
|
|
|
|
logger.debug(f'Created new node: {new_node.uuid}')
|
|
|
|
return extracted_nodes, node_episode_index_map
|
|
|
|
|
|
def _collapse_exact_duplicate_extracted_nodes(
|
|
extracted_nodes: list[EntityNode],
|
|
node_episode_index_map: dict[str, list[int]] | None = None,
|
|
) -> list[EntityNode]:
|
|
"""Collapse same-message duplicates with the same normalized name.
|
|
|
|
This is intentionally narrow: it only merges exact normalized-name duplicates that the
|
|
extraction prompt should already have emitted once. When duplicates disagree on specificity,
|
|
keep the more specific node (for example, `Person` over bare `Entity`).
|
|
|
|
When node_episode_index_map is provided, episode indices from discarded nodes are merged
|
|
into the canonical node's entry so attribution is preserved.
|
|
"""
|
|
if len(extracted_nodes) < 2:
|
|
return extracted_nodes
|
|
|
|
canonical_by_name: dict[str, EntityNode] = {}
|
|
ordered_names: list[str] = []
|
|
|
|
for node in extracted_nodes:
|
|
normalized_name = _normalize_string_exact(node.name)
|
|
existing = canonical_by_name.get(normalized_name)
|
|
if existing is None:
|
|
canonical_by_name[normalized_name] = node
|
|
ordered_names.append(normalized_name)
|
|
continue
|
|
|
|
existing_specific_labels = {label for label in existing.labels if label != 'Entity'}
|
|
node_specific_labels = {label for label in node.labels if label != 'Entity'}
|
|
if len(node_specific_labels) > len(existing_specific_labels) or (
|
|
len(node_specific_labels) == len(existing_specific_labels)
|
|
and len(node.name.strip()) > len(existing.name.strip())
|
|
):
|
|
old_canonical = existing
|
|
canonical_by_name[normalized_name] = node
|
|
# Merge episode indices: old canonical -> new canonical
|
|
if node_episode_index_map is not None:
|
|
old_indices = node_episode_index_map.pop(old_canonical.uuid, [])
|
|
new_indices = node_episode_index_map.get(node.uuid, [])
|
|
node_episode_index_map[node.uuid] = sorted(set(new_indices + old_indices))
|
|
elif node_episode_index_map is not None:
|
|
# Discard this node; merge its indices into the existing canonical
|
|
discarded_indices = node_episode_index_map.pop(node.uuid, [])
|
|
canonical_indices = node_episode_index_map.get(existing.uuid, [])
|
|
node_episode_index_map[existing.uuid] = sorted(
|
|
set(canonical_indices + discarded_indices)
|
|
)
|
|
|
|
return [canonical_by_name[name] for name in ordered_names]
|
|
|
|
|
|
def _merge_candidate_nodes(
|
|
candidate_nodes: list[EntityNode],
|
|
existing_nodes_override: list[EntityNode] | None,
|
|
) -> list[EntityNode]:
|
|
"""Deduplicate candidate nodes while preserving search order and overrides."""
|
|
merged_candidates = list(candidate_nodes)
|
|
if existing_nodes_override is not None:
|
|
merged_candidates.extend(existing_nodes_override)
|
|
|
|
seen_candidate_uuids: set[str] = set()
|
|
ordered_candidates: list[EntityNode] = []
|
|
for candidate in merged_candidates:
|
|
if candidate.uuid in seen_candidate_uuids:
|
|
continue
|
|
seen_candidate_uuids.add(candidate.uuid)
|
|
ordered_candidates.append(candidate)
|
|
|
|
return ordered_candidates
|
|
|
|
|
|
async def _collect_candidate_nodes(
|
|
clients: GraphitiClients,
|
|
extracted_nodes: list[EntityNode],
|
|
existing_nodes_override: list[EntityNode] | None,
|
|
) -> list[list[EntityNode]]:
|
|
"""Search per extracted name and return ordered candidates for each extracted node."""
|
|
search_results = await _semantic_candidate_search(clients, extracted_nodes)
|
|
|
|
return [_merge_candidate_nodes(result, existing_nodes_override) for result in search_results]
|
|
|
|
|
|
async def _semantic_candidate_search(
|
|
clients: GraphitiClients,
|
|
extracted_nodes: list[EntityNode],
|
|
) -> list[list[EntityNode]]:
|
|
"""Run direct cosine similarity search per extracted node without reranking."""
|
|
if not extracted_nodes:
|
|
return []
|
|
|
|
queries = [node.name.replace('\n', ' ') for node in extracted_nodes]
|
|
try:
|
|
query_vectors = await clients.embedder.create_batch(queries)
|
|
except NotImplementedError:
|
|
query_vectors = list(
|
|
await semaphore_gather(
|
|
*[clients.embedder.create(input_data=[query]) for query in queries]
|
|
)
|
|
)
|
|
|
|
return list(
|
|
await semaphore_gather(
|
|
*[
|
|
node_similarity_search(
|
|
clients.driver,
|
|
query_vector,
|
|
SearchFilters(),
|
|
[node.group_id],
|
|
NODE_DEDUP_CANDIDATE_LIMIT,
|
|
NODE_DEDUP_COSINE_MIN_SCORE,
|
|
)
|
|
for node, query_vector in zip(extracted_nodes, query_vectors, strict=True)
|
|
]
|
|
)
|
|
)
|
|
|
|
|
|
def _commit_resolution(
|
|
state: DedupResolutionState,
|
|
resolved_node: EntityNode | None,
|
|
uuid_map: dict[str, str],
|
|
duplicate_pairs: list[tuple[EntityNode, EntityNode]],
|
|
index: int,
|
|
) -> None:
|
|
"""Commit a single-node resolution result into the batch-level state."""
|
|
if resolved_node is not None:
|
|
state.resolved_nodes[index] = resolved_node
|
|
state.uuid_map.update(uuid_map)
|
|
state.duplicate_pairs.extend(duplicate_pairs)
|
|
|
|
|
|
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 {}
|
|
|
|
llm_extracted_nodes = [extracted_nodes[i] for i in state.unresolved_indices]
|
|
|
|
extracted_nodes_context = [
|
|
{
|
|
'id': i,
|
|
'name': node.name,
|
|
'entity_type': node.labels,
|
|
'entity_type_description': _get_entity_type_description(node.labels, entity_types_dict),
|
|
}
|
|
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 entity IDs: %s',
|
|
sample_size,
|
|
[ctx['id'] for ctx in extracted_nodes_context[:sample_size]],
|
|
)
|
|
if len(extracted_nodes_context) > 3:
|
|
logger.debug(
|
|
'Last %d entity IDs: %s',
|
|
sample_size,
|
|
[ctx['id'] for ctx in extracted_nodes_context[-sample_size:]],
|
|
)
|
|
|
|
existing_nodes_context = [
|
|
{
|
|
**candidate.attributes,
|
|
'candidate_id': i,
|
|
'name': candidate.name,
|
|
'entity_types': candidate.labels,
|
|
'summary': candidate.summary[:120] if candidate.summary else '',
|
|
}
|
|
for i, candidate in enumerate(indexes.existing_nodes)
|
|
]
|
|
|
|
# Build candidate_id -> node mapping for resolving duplicates by ID
|
|
candidates_by_id: dict[int, EntityNode] = {
|
|
i: node for i, node 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': (
|
|
[
|
|
{
|
|
'content': ep.content,
|
|
'timestamp': ep.valid_at.isoformat() if ep.valid_at else None,
|
|
}
|
|
for ep in previous_episodes
|
|
]
|
|
if previous_episodes is not None
|
|
else []
|
|
),
|
|
}
|
|
|
|
llm_response = await llm_client.generate_response(
|
|
prompt_library.dedupe_nodes.nodes(context),
|
|
response_model=NodeResolutions,
|
|
prompt_name='dedupe_nodes.nodes',
|
|
)
|
|
|
|
node_resolutions: list[NodeDuplicate] = NodeResolutions(**llm_response).entity_resolutions
|
|
|
|
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:
|
|
relative_id: int = resolution.id
|
|
duplicate_candidate_id: int = resolution.duplicate_candidate_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
|
|
|
|
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_candidate_id < 0:
|
|
resolved_node = extracted_node
|
|
elif duplicate_candidate_id in candidates_by_id:
|
|
resolved_node = _promote_resolved_node(
|
|
extracted_node, candidates_by_id[duplicate_candidate_id]
|
|
)
|
|
else:
|
|
logger.warning(
|
|
'Invalid duplicate_candidate_id %d for extracted node %s; treating as no duplicate.',
|
|
duplicate_candidate_id,
|
|
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]]]:
|
|
"""Resolve nodes with semantic retrieval first, then deterministic and LLM dedup."""
|
|
llm_client = clients.llm_client
|
|
candidate_nodes_by_extracted = await _collect_candidate_nodes(
|
|
clients,
|
|
extracted_nodes,
|
|
existing_nodes_override,
|
|
)
|
|
|
|
state = DedupResolutionState(
|
|
resolved_nodes=[None] * len(extracted_nodes),
|
|
uuid_map={},
|
|
unresolved_indices=[],
|
|
)
|
|
|
|
for idx, (node, candidates) in enumerate(
|
|
zip(extracted_nodes, candidate_nodes_by_extracted, strict=True)
|
|
):
|
|
if not candidates:
|
|
continue
|
|
|
|
indexes = _build_candidate_indexes(candidates)
|
|
local_state = DedupResolutionState(
|
|
resolved_nodes=[None], uuid_map={}, unresolved_indices=[]
|
|
)
|
|
_resolve_with_similarity([node], indexes, local_state)
|
|
if local_state.resolved_nodes[0] is not None:
|
|
_commit_resolution(
|
|
state,
|
|
local_state.resolved_nodes[0],
|
|
local_state.uuid_map,
|
|
local_state.duplicate_pairs,
|
|
idx,
|
|
)
|
|
continue
|
|
|
|
state.unresolved_indices.append(idx)
|
|
|
|
if state.unresolved_indices:
|
|
llm_candidate_nodes = _merge_candidate_nodes(
|
|
[
|
|
candidate
|
|
for idx in state.unresolved_indices
|
|
for candidate in candidate_nodes_by_extracted[idx]
|
|
],
|
|
None,
|
|
)
|
|
await _resolve_with_llm(
|
|
llm_client,
|
|
extracted_nodes,
|
|
_build_candidate_indexes(llm_candidate_nodes),
|
|
state,
|
|
episode,
|
|
previous_episodes,
|
|
entity_types,
|
|
)
|
|
|
|
if not state.unresolved_indices and not any(candidate_nodes_by_extracted):
|
|
logger.debug('No semantic dedup candidates found; keeping all extracted nodes as new')
|
|
|
|
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(
|
|
'Resolved nodes: %s',
|
|
[node.uuid for node in state.resolved_nodes if node is not None],
|
|
)
|
|
|
|
return (
|
|
[node for node in state.resolved_nodes if node is not None],
|
|
state.uuid_map,
|
|
state.duplicate_pairs,
|
|
)
|
|
|
|
|
|
def _build_edges_by_node(edges: list[EntityEdge] | None) -> dict[str, list[EntityEdge]]:
|
|
"""Build a dictionary mapping node UUIDs to their connected edges."""
|
|
edges_by_node: dict[str, list[EntityEdge]] = {}
|
|
if not edges:
|
|
return edges_by_node
|
|
for edge in edges:
|
|
if edge.source_node_uuid not in edges_by_node:
|
|
edges_by_node[edge.source_node_uuid] = []
|
|
if edge.target_node_uuid not in edges_by_node:
|
|
edges_by_node[edge.target_node_uuid] = []
|
|
edges_by_node[edge.source_node_uuid].append(edge)
|
|
edges_by_node[edge.target_node_uuid].append(edge)
|
|
return edges_by_node
|
|
|
|
|
|
async def extract_attributes_from_nodes(
|
|
clients: GraphitiClients,
|
|
nodes: list[EntityNode],
|
|
episode: EpisodicNode | list[EpisodicNode] | None = None,
|
|
previous_episodes: list[EpisodicNode] | None = None,
|
|
entity_types: dict[str, type[BaseModel]] | None = None,
|
|
should_summarize_node: NodeSummaryFilter | None = None,
|
|
edges: list[EntityEdge] | None = None,
|
|
skip_fact_appending: bool = False,
|
|
include_type_descriptions: bool = False,
|
|
) -> list[EntityNode]:
|
|
llm_client = clients.llm_client
|
|
embedder = clients.embedder
|
|
|
|
# Pre-build edges lookup for O(E + N) instead of O(N * E)
|
|
edges_by_node = _build_edges_by_node(edges)
|
|
|
|
# Extract attributes in parallel (per-entity calls)
|
|
attribute_results: list[dict[str, Any]] = await semaphore_gather(
|
|
*[
|
|
_extract_entity_attributes(
|
|
llm_client,
|
|
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
|
|
),
|
|
)
|
|
for node in nodes
|
|
]
|
|
)
|
|
|
|
# _extract_entity_attributes returns the already-merged attribute dict
|
|
# (overlay of prior + cap-kept fields), so direct assignment is the merge.
|
|
for node, attributes in zip(nodes, attribute_results, strict=True):
|
|
node.attributes = attributes
|
|
|
|
# Extract summaries in batch
|
|
await _extract_entity_summaries_batch(
|
|
llm_client,
|
|
nodes,
|
|
episode,
|
|
previous_episodes,
|
|
should_summarize_node,
|
|
edges_by_node,
|
|
skip_fact_appending=skip_fact_appending,
|
|
entity_types=entity_types if include_type_descriptions else None,
|
|
)
|
|
|
|
await create_entity_node_embeddings(embedder, nodes)
|
|
|
|
return nodes
|
|
|
|
|
|
async def _extract_entity_attributes(
|
|
llm_client: LLMClient,
|
|
node: EntityNode,
|
|
episode: EpisodicNode | list[EpisodicNode] | None,
|
|
previous_episodes: list[EpisodicNode] | None,
|
|
entity_type: type[BaseModel] | None,
|
|
) -> dict[str, Any]:
|
|
if entity_type is None or len(entity_type.model_fields) == 0:
|
|
return {}
|
|
|
|
attributes_context = _build_episode_context(
|
|
# should not include summary
|
|
node_data={
|
|
'name': node.name,
|
|
'entity_types': node.labels,
|
|
'attributes': node.attributes,
|
|
},
|
|
episode=episode,
|
|
previous_episodes=previous_episodes,
|
|
)
|
|
|
|
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,
|
|
prompt_name='extract_nodes.extract_attributes',
|
|
attribute_extraction=True,
|
|
)
|
|
|
|
# Overlay merge: cap-dropped or LLM-omitted fields keep prior values.
|
|
# See attribute_utils for the merge_mode contract; the edge path uses 'replace'.
|
|
merged, _ = apply_capped_attributes(
|
|
llm_response,
|
|
entity_type,
|
|
node.attributes,
|
|
merge_mode='overlay',
|
|
prompt_name='extract_nodes.extract_attributes',
|
|
entity_uuid=node.uuid,
|
|
group_id=node.group_id,
|
|
)
|
|
|
|
# Shape validation only — we discard the validated instance because returning
|
|
# `model_dump()` would expand defaults across all fields and clobber prior
|
|
# values that the merge above just preserved.
|
|
entity_type(**merged)
|
|
|
|
return merged
|
|
|
|
|
|
async def _extract_entity_summaries_batch(
|
|
llm_client: LLMClient,
|
|
nodes: list[EntityNode],
|
|
episode: EpisodicNode | list[EpisodicNode] | None,
|
|
previous_episodes: list[EpisodicNode] | None,
|
|
should_summarize_node: NodeSummaryFilter | None,
|
|
edges_by_node: dict[str, list[EntityEdge]],
|
|
*,
|
|
skip_fact_appending: bool = False,
|
|
entity_types: dict[str, type[BaseModel]] | None = None,
|
|
) -> None:
|
|
"""Extract summaries for multiple entities in batched LLM calls.
|
|
|
|
When skip_fact_appending is False (default), nodes with short summaries get edge
|
|
facts appended directly without an LLM call. Nodes needing summarization are
|
|
partitioned into flights of MAX_NODES and processed with separate LLM calls.
|
|
|
|
When skip_fact_appending is True, the raw fact-append shortcut is bypassed and all
|
|
nodes are routed through LLM summarization using an episode-based prompt that
|
|
matches the async graph summary worker.
|
|
"""
|
|
# Determine which nodes need LLM summarization vs direct edge fact appending
|
|
nodes_needing_llm: list[EntityNode] = []
|
|
|
|
for node in nodes:
|
|
# Check if node should be summarized at all
|
|
if should_summarize_node is not None and not await should_summarize_node(node):
|
|
continue
|
|
|
|
if skip_fact_appending:
|
|
# Always route through LLM — no raw fact concatenation.
|
|
if episode is not None or node.summary:
|
|
nodes_needing_llm.append(node)
|
|
continue
|
|
|
|
node_edges = edges_by_node.get(node.uuid, [])
|
|
|
|
# Build summary with edge facts appended
|
|
summary_with_edges = node.summary
|
|
if node_edges:
|
|
edge_facts = '\n'.join(edge.fact for edge in node_edges if edge.fact)
|
|
summary_with_edges = f'{summary_with_edges}\n{edge_facts}'.strip()
|
|
|
|
# If summary is close to the persisted limit, use it directly (append edge facts, no LLM call)
|
|
if summary_with_edges and len(summary_with_edges) <= MAX_SUMMARY_CHARS * 2:
|
|
node.summary = summary_with_edges
|
|
continue
|
|
|
|
# Skip if no summary content and no episode to generate from
|
|
if not summary_with_edges and episode is None:
|
|
continue
|
|
|
|
# This node needs LLM summarization
|
|
nodes_needing_llm.append(node)
|
|
|
|
# If no nodes need LLM summarization, return early
|
|
if not nodes_needing_llm:
|
|
return
|
|
|
|
# Partition nodes into flights of MAX_NODES
|
|
node_flights = [
|
|
nodes_needing_llm[i : i + MAX_NODES] for i in range(0, len(nodes_needing_llm), MAX_NODES)
|
|
]
|
|
|
|
# Process flights in parallel
|
|
await semaphore_gather(
|
|
*[
|
|
_process_summary_flight(
|
|
llm_client,
|
|
flight,
|
|
episode,
|
|
previous_episodes,
|
|
use_episode_prompt=skip_fact_appending,
|
|
entity_types=entity_types,
|
|
)
|
|
for flight in node_flights
|
|
]
|
|
)
|
|
|
|
|
|
async def _process_summary_flight(
|
|
llm_client: LLMClient,
|
|
nodes: list[EntityNode],
|
|
episode: EpisodicNode | list[EpisodicNode] | None,
|
|
previous_episodes: list[EpisodicNode] | None,
|
|
*,
|
|
use_episode_prompt: bool = False,
|
|
entity_types: dict[str, type[BaseModel]] | None = None,
|
|
) -> None:
|
|
"""Process a single flight of nodes for batch summarization."""
|
|
# Build entity type descriptions from docstrings, stripping GOOD/BAD
|
|
# few-shot examples that are intended for extraction prompts only.
|
|
entity_type_descriptions: dict[str, str] = {}
|
|
if entity_types is not None:
|
|
for type_name, type_model in entity_types.items():
|
|
if type_model.__doc__:
|
|
entity_type_descriptions[type_name] = _truncate_type_description(type_model.__doc__)
|
|
|
|
# Build context for batch summarization
|
|
entities_context = [
|
|
{
|
|
'name': node.name,
|
|
'summary': node.summary,
|
|
'entity_types': node.labels,
|
|
'attributes': node.attributes,
|
|
}
|
|
for node in nodes
|
|
]
|
|
|
|
if episode is None:
|
|
episode_content = ''
|
|
elif isinstance(episode, list):
|
|
episode_content = concatenate_episodes(episode)
|
|
else:
|
|
episode_content = episode.content
|
|
|
|
batch_context: dict[str, Any] = {
|
|
'entities': entities_context,
|
|
'episode_content': episode_content,
|
|
'previous_episodes': (
|
|
[
|
|
{
|
|
'content': ep.content,
|
|
'timestamp': ep.valid_at.isoformat() if ep.valid_at else None,
|
|
}
|
|
for ep in previous_episodes
|
|
]
|
|
if previous_episodes is not None
|
|
else []
|
|
),
|
|
'entity_type_descriptions': entity_type_descriptions,
|
|
}
|
|
|
|
# Get group_id from the first node (all nodes in a batch should have same group_id)
|
|
group_id = nodes[0].group_id if nodes else None
|
|
|
|
if use_episode_prompt:
|
|
prompt = prompt_library.extract_nodes.extract_entity_summaries_from_episodes(batch_context)
|
|
prompt_name = 'extract_nodes.extract_entity_summaries_from_episodes'
|
|
else:
|
|
prompt = prompt_library.extract_nodes.extract_summaries_batch(batch_context)
|
|
prompt_name = 'extract_nodes.extract_summaries_batch'
|
|
|
|
llm_response = await llm_client.generate_response(
|
|
prompt,
|
|
response_model=SummarizedEntities,
|
|
model_size=ModelSize.small,
|
|
group_id=group_id,
|
|
prompt_name=prompt_name,
|
|
)
|
|
|
|
# Build case-insensitive name -> nodes mapping (handles duplicates)
|
|
name_to_nodes: dict[str, list[EntityNode]] = {}
|
|
for node in nodes:
|
|
key = node.name.lower()
|
|
if key not in name_to_nodes:
|
|
name_to_nodes[key] = []
|
|
name_to_nodes[key].append(node)
|
|
|
|
# Apply summaries from LLM response
|
|
summaries_response = SummarizedEntities(**llm_response)
|
|
for summarized_entity in summaries_response.summaries:
|
|
matching_nodes = name_to_nodes.get(summarized_entity.name.lower(), [])
|
|
if matching_nodes:
|
|
truncated_summary = truncate_at_sentence(summarized_entity.summary, MAX_SUMMARY_CHARS)
|
|
for node in matching_nodes:
|
|
node.summary = truncated_summary
|
|
else:
|
|
logger.warning(
|
|
'LLM returned summary for unknown entity (first 30 chars): %.30s',
|
|
summarized_entity.name,
|
|
)
|
|
|
|
|
|
def _build_episode_context(
|
|
node_data: dict[str, Any],
|
|
episode: EpisodicNode | list[EpisodicNode] | None,
|
|
previous_episodes: list[EpisodicNode] | None,
|
|
) -> dict[str, Any]:
|
|
if episode is None:
|
|
episode_content = ''
|
|
elif isinstance(episode, list):
|
|
episode_content = concatenate_episodes(episode)
|
|
else:
|
|
episode_content = episode.content
|
|
return {
|
|
'node': node_data,
|
|
'episode_content': episode_content,
|
|
'previous_episodes': (
|
|
[
|
|
{
|
|
'content': ep.content,
|
|
'timestamp': ep.valid_at.isoformat() if ep.valid_at else None,
|
|
}
|
|
for ep in previous_episodes
|
|
]
|
|
if previous_episodes is not None
|
|
else []
|
|
),
|
|
}
|