graphiti-core 0.21.0rc5__py3-none-any.whl → 0.21.0rc7__py3-none-any.whl

graphiti_core/graphiti.py

@@ -79,7 +79,6 @@ from graphiti_core.utils.maintenance.community_operations import (
     update_community,
 )
 from graphiti_core.utils.maintenance.edge_operations import (
-    build_duplicate_of_edges,
     build_episodic_edges,
     extract_edges,
     resolve_extracted_edge,
@@ -503,7 +502,7 @@ class Graphiti:
             )
 
             # Extract edges and resolve nodes
-            (nodes, uuid_map, node_duplicates), extracted_edges = await semaphore_gather(
+            (nodes, uuid_map, _), extracted_edges = await semaphore_gather(
                 resolve_extracted_nodes(
                     self.clients,
                     extracted_nodes,
@@ -540,9 +539,7 @@ class Graphiti:
                 max_coroutines=self.max_coroutines,
             )
 
-            duplicate_of_edges = build_duplicate_of_edges(episode, now, node_duplicates)
-
-            entity_edges = resolved_edges + invalidated_edges + duplicate_of_edges
+            entity_edges = resolved_edges + invalidated_edges
 
             episodic_edges = build_episodic_edges(nodes, episode.uuid, now)
 
@@ -1073,6 +1070,7 @@ class Graphiti:
                 group_id=edge.group_id,
             ),
             None,
+            None,
             self.ensure_ascii,
         )
 

graphiti_core/llm_client/client.py

@@ -32,9 +32,19 @@ from .errors import RateLimitError
 DEFAULT_TEMPERATURE = 0
 DEFAULT_CACHE_DIR = './llm_cache'
 
-MULTILINGUAL_EXTRACTION_RESPONSES = (
-    '\n\nAny extracted information should be returned in the same language as it was written in.'
-)
+
+def get_extraction_language_instruction() -> str:
+    """Returns instruction for language extraction behavior.
+
+    Override this function to customize language extraction:
+    - Return empty string to disable multilingual instructions
+    - Return custom instructions for specific language requirements
+
+    Returns:
+        str: Language instruction to append to system messages
+    """
+    return '\n\nAny extracted information should be returned in the same language as it was written in.'
+
 
 logger = logging.getLogger(__name__)
 
@@ -145,7 +155,7 @@ class LLMClient(ABC):
             )
 
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+        messages[0].content += get_extraction_language_instruction()
 
         if self.cache_enabled and self.cache_dir is not None:
             cache_key = self._get_cache_key(messages)

graphiti_core/llm_client/gemini_client.py

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, ClassVar
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import LLMConfig, ModelSize
 from .errors import RateLimitError
 
@@ -376,7 +376,7 @@ class GeminiClient(LLMClient):
         last_output = None
 
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+        messages[0].content += get_extraction_language_instruction()
 
         while retry_count < self.MAX_RETRIES:
             try:

graphiti_core/llm_client/openai_base_client.py

@@ -25,7 +25,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError, RefusalError
 
@@ -184,7 +184,7 @@ class BaseOpenAIClient(LLMClient):
         last_error = None
 
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+        messages[0].content += get_extraction_language_instruction()
 
         while retry_count <= self.MAX_RETRIES:
             try:

graphiti_core/llm_client/openai_generic_client.py

@@ -25,7 +25,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError, RefusalError
 
@@ -136,7 +136,7 @@ class OpenAIGenericClient(LLMClient):
             )
 
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+        messages[0].content += get_extraction_language_instruction()
 
         while retry_count <= self.MAX_RETRIES:
             try:

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/prompts/extract_nodes.py

@@ -152,7 +152,8 @@ Indicate the classified entity type by providing its entity_type_id.
 
 Guidelines:
 1. Always try to extract an entities that the JSON represents. This will often be something like a "name" or "user field
-2. Do NOT extract any properties that contain dates
+2. Extract all entities mentioned in all other properties throughout the JSON structure
+3. Do NOT extract any properties that contain dates
 """
     return [
         Message(role='system', content=sys_prompt),

graphiti_core/utils/bulk_utils.py

@@ -43,8 +43,14 @@ from graphiti_core.models.nodes.node_db_queries import (
     get_entity_node_save_bulk_query,
     get_episode_node_save_bulk_query,
 )
-from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
+from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.utils.datetime_utils import convert_datetimes_to_strings
+from graphiti_core.utils.maintenance.dedup_helpers import (
+    DedupResolutionState,
+    _build_candidate_indexes,
+    _normalize_string_exact,
+    _resolve_with_similarity,
+)
 from graphiti_core.utils.maintenance.edge_operations import (
     extract_edges,
     resolve_extracted_edge,
@@ -63,6 +69,38 @@ logger = logging.getLogger(__name__)
 CHUNK_SIZE = 10
 
 
+def _build_directed_uuid_map(pairs: list[tuple[str, str]]) -> dict[str, str]:
+    """Collapse alias -> canonical chains while preserving direction.
+
+    The incoming pairs represent directed mappings discovered during node dedupe. We use a simple
+    union-find with iterative path compression to ensure every source UUID resolves to its ultimate
+    canonical target, even if aliases appear lexicographically smaller than the canonical UUID.
+    """
+
+    parent: dict[str, str] = {}
+
+    def find(uuid: str) -> str:
+        """Directed union-find lookup using iterative path compression."""
+        parent.setdefault(uuid, uuid)
+        root = uuid
+        while parent[root] != root:
+            root = parent[root]
+
+        while parent[uuid] != root:
+            next_uuid = parent[uuid]
+            parent[uuid] = root
+            uuid = next_uuid
+
+        return root
+
+    for source_uuid, target_uuid in pairs:
+        parent.setdefault(source_uuid, source_uuid)
+        parent.setdefault(target_uuid, target_uuid)
+        parent[find(source_uuid)] = find(target_uuid)
+
+    return {uuid: find(uuid) for uuid in parent}
+
+
 class RawEpisode(BaseModel):
     name: str
     uuid: str | None = Field(default=None)
@@ -266,83 +304,111 @@ async def dedupe_nodes_bulk(
     episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]],
     entity_types: dict[str, type[BaseModel]] | None = None,
 ) -> tuple[dict[str, list[EntityNode]], dict[str, str]]:
-    embedder = clients.embedder
-    min_score = 0.8
-
-    # generate embeddings
-    await semaphore_gather(
-        *[create_entity_node_embeddings(embedder, nodes) for nodes in extracted_nodes]
-    )
-
-    # Find similar results
-    dedupe_tuples: list[tuple[list[EntityNode], list[EntityNode]]] = []
-    for i, nodes_i in enumerate(extracted_nodes):
-        existing_nodes: list[EntityNode] = []
-        for j, nodes_j in enumerate(extracted_nodes):
-            if i == j:
-                continue
-            existing_nodes += nodes_j
-
-        candidates_i: list[EntityNode] = []
-        for node in nodes_i:
-            for existing_node in existing_nodes:
-                # Approximate BM25 by checking for word overlaps (this is faster than creating many in-memory indices)
-                # This approach will cast a wider net than BM25, which is ideal for this use case
-                node_words = set(node.name.lower().split())
-                existing_node_words = set(existing_node.name.lower().split())
-                has_overlap = not node_words.isdisjoint(existing_node_words)
-                if has_overlap:
-                    candidates_i.append(existing_node)
-                    continue
+    """Resolve entity duplicates across an in-memory batch using a two-pass strategy.
 
-                # Check for semantic similarity even if there is no overlap
-                similarity = np.dot(
-                    normalize_l2(node.name_embedding or []),
-                    normalize_l2(existing_node.name_embedding or []),
-                )
-                if similarity >= min_score:
-                    candidates_i.append(existing_node)
-
-        dedupe_tuples.append((nodes_i, candidates_i))
+    1. Run :func:`resolve_extracted_nodes` for every episode in parallel so each batch item is
+       reconciled against the live graph just like the non-batch flow.
+    2. Re-run the deterministic similarity heuristics across the union of resolved nodes to catch
+       duplicates that only co-occur inside this batch, emitting a canonical UUID map that callers
+       can apply to edges and persistence.
+    """
 
-    # Determine Node Resolutions
-    bulk_node_resolutions: list[
-        tuple[list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]]]
-    ] = await semaphore_gather(
+    first_pass_results = await semaphore_gather(
         *[
             resolve_extracted_nodes(
                 clients,
-                dedupe_tuple[0],
+                nodes,
                 episode_tuples[i][0],
                 episode_tuples[i][1],
                 entity_types,
-                existing_nodes_override=dedupe_tuples[i][1],
             )
-            for i, dedupe_tuple in enumerate(dedupe_tuples)
+            for i, nodes in enumerate(extracted_nodes)
         ]
     )
 
-    # Collect all duplicate pairs sorted by uuid
+    episode_resolutions: list[tuple[str, list[EntityNode]]] = []
+    per_episode_uuid_maps: list[dict[str, str]] = []
     duplicate_pairs: list[tuple[str, str]] = []
-    for _, _, duplicates in bulk_node_resolutions:
-        for duplicate in duplicates:
-            n, m = duplicate
-            duplicate_pairs.append((n.uuid, m.uuid))
 
-    # Now we compress the duplicate_map, so that 3 -> 2 and 2 -> becomes 3 -> 1 (sorted by uuid)
-    compressed_map: dict[str, str] = compress_uuid_map(duplicate_pairs)
+    for (resolved_nodes, uuid_map, duplicates), (episode, _) in zip(
+        first_pass_results, episode_tuples, strict=True
+    ):
+        episode_resolutions.append((episode.uuid, resolved_nodes))
+        per_episode_uuid_maps.append(uuid_map)
+        duplicate_pairs.extend((source.uuid, target.uuid) for source, target in duplicates)
+
+    canonical_nodes: dict[str, EntityNode] = {}
+    for _, resolved_nodes in episode_resolutions:
+        for node in resolved_nodes:
+            # NOTE: this loop is O(n^2) in the number of nodes inside the batch because we rebuild
+            # the MinHash index for the accumulated canonical pool each time. The LRU-backed
+            # shingle cache keeps the constant factors low for typical batch sizes (≤ CHUNK_SIZE),
+            # but if batches grow significantly we should switch to an incremental index or chunked
+            # processing.
+            if not canonical_nodes:
+                canonical_nodes[node.uuid] = node
+                continue
 
-    node_uuid_map: dict[str, EntityNode] = {
-        node.uuid: node for nodes in extracted_nodes for node in nodes
-    }
+            existing_candidates = list(canonical_nodes.values())
+            normalized = _normalize_string_exact(node.name)
+            exact_match = next(
+                (
+                    candidate
+                    for candidate in existing_candidates
+                    if _normalize_string_exact(candidate.name) == normalized
+                ),
+                None,
+            )
+            if exact_match is not None:
+                if exact_match.uuid != node.uuid:
+                    duplicate_pairs.append((node.uuid, exact_match.uuid))
+                continue
+
+            indexes = _build_candidate_indexes(existing_candidates)
+            state = DedupResolutionState(
+                resolved_nodes=[None],
+                uuid_map={},
+                unresolved_indices=[],
+            )
+            _resolve_with_similarity([node], indexes, state)
+
+            resolved = state.resolved_nodes[0]
+            if resolved is None:
+                canonical_nodes[node.uuid] = node
+                continue
+
+            canonical_uuid = resolved.uuid
+            canonical_nodes.setdefault(canonical_uuid, resolved)
+            if canonical_uuid != node.uuid:
+                duplicate_pairs.append((node.uuid, canonical_uuid))
+
+    union_pairs: list[tuple[str, str]] = []
+    for uuid_map in per_episode_uuid_maps:
+        union_pairs.extend(uuid_map.items())
+    union_pairs.extend(duplicate_pairs)
+
+    compressed_map: dict[str, str] = _build_directed_uuid_map(union_pairs)
 
     nodes_by_episode: dict[str, list[EntityNode]] = {}
-    for i, nodes in enumerate(extracted_nodes):
-        episode = episode_tuples[i][0]
+    for episode_uuid, resolved_nodes in episode_resolutions:
+        deduped_nodes: list[EntityNode] = []
+        seen: set[str] = set()
+        for node in resolved_nodes:
+            canonical_uuid = compressed_map.get(node.uuid, node.uuid)
+            if canonical_uuid in seen:
+                continue
+            seen.add(canonical_uuid)
+            canonical_node = canonical_nodes.get(canonical_uuid)
+            if canonical_node is None:
+                logger.error(
+                    'Canonical node %s missing during batch dedupe; falling back to %s',
+                    canonical_uuid,
+                    node.uuid,
+                )
+                canonical_node = node
+            deduped_nodes.append(canonical_node)
 
-        nodes_by_episode[episode.uuid] = [
-            node_uuid_map[compressed_map.get(node.uuid, node.uuid)] for node in nodes
-        ]
+        nodes_by_episode[episode_uuid] = deduped_nodes
 
     return nodes_by_episode, compressed_map
 
@@ -411,6 +477,7 @@ async def dedupe_edges_bulk(
                 candidates,
                 episode,
                 edge_types,
+                set(edge_types),
                 clients.ensure_ascii,
             )
             for episode, edge, candidates in dedupe_tuples