graphiti-core 0.17.4__py3-none-any.whl → 0.25.3__py3-none-any.whl

graphiti_core/tracer.py

@@ -0,0 +1,193 @@
+"""
+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.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+from contextlib import AbstractContextManager, contextmanager, suppress
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span, StatusCode
+
+try:
+    from opentelemetry.trace import Span, StatusCode
+
+    OTEL_AVAILABLE = True
+except ImportError:
+    OTEL_AVAILABLE = False
+
+
+class TracerSpan(ABC):
+    """Abstract base class for tracer spans."""
+
+    @abstractmethod
+    def add_attributes(self, attributes: dict[str, Any]) -> None:
+        """Add attributes to the span."""
+        pass
+
+    @abstractmethod
+    def set_status(self, status: str, description: str | None = None) -> None:
+        """Set the status of the span."""
+        pass
+
+    @abstractmethod
+    def record_exception(self, exception: Exception) -> None:
+        """Record an exception in the span."""
+        pass
+
+
+class Tracer(ABC):
+    """Abstract base class for tracers."""
+
+    @abstractmethod
+    def start_span(self, name: str) -> AbstractContextManager[TracerSpan]:
+        """Start a new span with the given name."""
+        pass
+
+
+class NoOpSpan(TracerSpan):
+    """No-op span implementation that does nothing."""
+
+    def add_attributes(self, attributes: dict[str, Any]) -> None:
+        pass
+
+    def set_status(self, status: str, description: str | None = None) -> None:
+        pass
+
+    def record_exception(self, exception: Exception) -> None:
+        pass
+
+
+class NoOpTracer(Tracer):
+    """No-op tracer implementation that does nothing."""
+
+    @contextmanager
+    def start_span(self, name: str) -> Generator[NoOpSpan, None, None]:
+        """Return a no-op span."""
+        yield NoOpSpan()
+
+
+class OpenTelemetrySpan(TracerSpan):
+    """Wrapper for OpenTelemetry span."""
+
+    def __init__(self, span: 'Span'):
+        self._span = span
+
+    def add_attributes(self, attributes: dict[str, Any]) -> None:
+        """Add attributes to the OpenTelemetry span."""
+        try:
+            # Filter out None values and convert all values to appropriate types
+            filtered_attrs = {}
+            for key, value in attributes.items():
+                if value is not None:
+                    # Convert to string if not a primitive type
+                    if isinstance(value, str | int | float | bool):
+                        filtered_attrs[key] = value
+                    else:
+                        filtered_attrs[key] = str(value)
+
+            if filtered_attrs:
+                self._span.set_attributes(filtered_attrs)
+        except Exception:
+            # Silently ignore tracing errors
+            pass
+
+    def set_status(self, status: str, description: str | None = None) -> None:
+        """Set the status of the OpenTelemetry span."""
+        try:
+            if OTEL_AVAILABLE:
+                if status == 'error':
+                    self._span.set_status(StatusCode.ERROR, description)
+                elif status == 'ok':
+                    self._span.set_status(StatusCode.OK, description)
+        except Exception:
+            # Silently ignore tracing errors
+            pass
+
+    def record_exception(self, exception: Exception) -> None:
+        """Record an exception in the OpenTelemetry span."""
+        with suppress(Exception):
+            self._span.record_exception(exception)
+
+
+class OpenTelemetryTracer(Tracer):
+    """Wrapper for OpenTelemetry tracer with configurable span name prefix."""
+
+    def __init__(self, tracer: Any, span_prefix: str = 'graphiti'):
+        """
+        Initialize the OpenTelemetry tracer wrapper.
+
+        Parameters
+        ----------
+        tracer : opentelemetry.trace.Tracer
+            The OpenTelemetry tracer instance.
+        span_prefix : str, optional
+            Prefix to prepend to all span names. Defaults to 'graphiti'.
+        """
+        if not OTEL_AVAILABLE:
+            raise ImportError(
+                'OpenTelemetry is not installed. Install it with: pip install opentelemetry-api'
+            )
+        self._tracer = tracer
+        self._span_prefix = span_prefix.rstrip('.')
+
+    @contextmanager
+    def start_span(self, name: str) -> Generator[OpenTelemetrySpan | NoOpSpan, None, None]:
+        """Start a new OpenTelemetry span with the configured prefix."""
+        try:
+            full_name = f'{self._span_prefix}.{name}'
+            with self._tracer.start_as_current_span(full_name) as span:
+                yield OpenTelemetrySpan(span)
+        except Exception:
+            # If tracing fails, yield a no-op span to prevent breaking the operation
+            yield NoOpSpan()
+
+
+def create_tracer(otel_tracer: Any | None = None, span_prefix: str = 'graphiti') -> Tracer:
+    """
+    Create a tracer instance.
+
+    Parameters
+    ----------
+    otel_tracer : opentelemetry.trace.Tracer | None, optional
+        An OpenTelemetry tracer instance. If None, a no-op tracer is returned.
+    span_prefix : str, optional
+        Prefix to prepend to all span names. Defaults to 'graphiti'.
+
+    Returns
+    -------
+    Tracer
+        A tracer instance (either OpenTelemetryTracer or NoOpTracer).
+
+    Examples
+    --------
+    Using with OpenTelemetry:
+
+    >>> from opentelemetry import trace
+    >>> otel_tracer = trace.get_tracer(__name__)
+    >>> tracer = create_tracer(otel_tracer, span_prefix='myapp.graphiti')
+
+    Using no-op tracer:
+
+    >>> tracer = create_tracer()  # Returns NoOpTracer
+    """
+    if otel_tracer is None:
+        return NoOpTracer()
+
+    if not OTEL_AVAILABLE:
+        return NoOpTracer()
+
+    return OpenTelemetryTracer(otel_tracer, span_prefix)

graphiti_core/utils/bulk_utils.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import json
 import logging
 import typing
 from datetime import datetime
@@ -22,22 +23,31 @@ import numpy as np
 from pydantic import BaseModel, Field
 from typing_extensions import Any
 
-from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
+from graphiti_core.driver.driver import (
+    GraphDriver,
+    GraphDriverSession,
+    GraphProvider,
+)
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge, create_entity_edge_embeddings
 from graphiti_core.embedder import EmbedderClient
-from graphiti_core.graph_queries import (
-    get_entity_edge_save_bulk_query,
-    get_entity_node_save_bulk_query,
-)
 from graphiti_core.graphiti_types import GraphitiClients
 from graphiti_core.helpers import normalize_l2, semaphore_gather
 from graphiti_core.models.edges.edge_db_queries import (
-    EPISODIC_EDGE_SAVE_BULK,
+    get_entity_edge_save_bulk_query,
+    get_episodic_edge_save_bulk_query,
 )
 from graphiti_core.models.nodes.node_db_queries import (
-    EPISODIC_NODE_SAVE_BULK,
+    get_entity_node_save_bulk_query,
+    get_episode_node_save_bulk_query,
+)
+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.nodes import EntityNode, EpisodeType, EpisodicNode, create_entity_node_embeddings
 from graphiti_core.utils.maintenance.edge_operations import (
     extract_edges,
     resolve_extracted_edge,
@@ -56,6 +66,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)
@@ -118,24 +160,33 @@ async def add_nodes_and_edges_bulk_tx(
     episodes = [dict(episode) for episode in episodic_nodes]
     for episode in episodes:
         episode['source'] = str(episode['source'].value)
-    nodes: list[dict[str, Any]] = []
+        episode.pop('labels', None)
+
+    nodes = []
+
     for node in entity_nodes:
         if node.name_embedding is None:
             await node.generate_name_embedding(embedder)
+
         entity_data: dict[str, Any] = {
             'uuid': node.uuid,
             'name': node.name,
-            'name_embedding': node.name_embedding,
             'group_id': node.group_id,
             'summary': node.summary,
             'created_at': node.created_at,
+            'name_embedding': node.name_embedding,
+            'labels': list(set(node.labels + ['Entity'])),
         }
 
-        entity_data.update(node.attributes or {})
-        entity_data['labels'] = list(set(node.labels + ['Entity']))
+        if driver.provider == GraphProvider.KUZU:
+            attributes = convert_datetimes_to_strings(node.attributes) if node.attributes else {}
+            entity_data['attributes'] = json.dumps(attributes)
+        else:
+            entity_data.update(node.attributes or {})
+
         nodes.append(entity_data)
 
-    edges: list[dict[str, Any]] = []
+    edges = []
     for edge in entity_edges:
         if edge.fact_embedding is None:
             await edge.generate_embedding(embedder)
@@ -145,35 +196,68 @@ async def add_nodes_and_edges_bulk_tx(
             'target_node_uuid': edge.target_node_uuid,
             'name': edge.name,
             'fact': edge.fact,
-            'fact_embedding': edge.fact_embedding,
             'group_id': edge.group_id,
             'episodes': edge.episodes,
             'created_at': edge.created_at,
             'expired_at': edge.expired_at,
             'valid_at': edge.valid_at,
             'invalid_at': edge.invalid_at,
+            'fact_embedding': edge.fact_embedding,
         }
 
-        edge_data.update(edge.attributes or {})
+        if driver.provider == GraphProvider.KUZU:
+            attributes = convert_datetimes_to_strings(edge.attributes) if edge.attributes else {}
+            edge_data['attributes'] = json.dumps(attributes)
+        else:
+            edge_data.update(edge.attributes or {})
+
         edges.append(edge_data)
 
-    await tx.run(EPISODIC_NODE_SAVE_BULK, episodes=episodes)
-    entity_node_save_bulk = get_entity_node_save_bulk_query(nodes, driver.provider)
-    await tx.run(entity_node_save_bulk, nodes=nodes)
-    await tx.run(
-        EPISODIC_EDGE_SAVE_BULK, episodic_edges=[edge.model_dump() for edge in episodic_edges]
-    )
-    entity_edge_save_bulk = get_entity_edge_save_bulk_query(driver.provider)
-    await tx.run(entity_edge_save_bulk, entity_edges=edges)
+    if driver.graph_operations_interface:
+        await driver.graph_operations_interface.episodic_node_save_bulk(None, driver, tx, episodes)
+        await driver.graph_operations_interface.node_save_bulk(None, driver, tx, nodes)
+        await driver.graph_operations_interface.episodic_edge_save_bulk(
+            None, driver, tx, [edge.model_dump() for edge in episodic_edges]
+        )
+        await driver.graph_operations_interface.edge_save_bulk(None, driver, tx, edges)
+
+    elif driver.provider == GraphProvider.KUZU:
+        # FIXME: Kuzu's UNWIND does not currently support STRUCT[] type properly, so we insert the data one by one instead for now.
+        episode_query = get_episode_node_save_bulk_query(driver.provider)
+        for episode in episodes:
+            await tx.run(episode_query, **episode)
+        entity_node_query = get_entity_node_save_bulk_query(driver.provider, nodes)
+        for node in nodes:
+            await tx.run(entity_node_query, **node)
+        entity_edge_query = get_entity_edge_save_bulk_query(driver.provider)
+        for edge in edges:
+            await tx.run(entity_edge_query, **edge)
+        episodic_edge_query = get_episodic_edge_save_bulk_query(driver.provider)
+        for edge in episodic_edges:
+            await tx.run(episodic_edge_query, **edge.model_dump())
+    else:
+        await tx.run(get_episode_node_save_bulk_query(driver.provider), episodes=episodes)
+        await tx.run(
+            get_entity_node_save_bulk_query(driver.provider, nodes),
+            nodes=nodes,
+        )
+        await tx.run(
+            get_episodic_edge_save_bulk_query(driver.provider),
+            episodic_edges=[edge.model_dump() for edge in episodic_edges],
+        )
+        await tx.run(
+            get_entity_edge_save_bulk_query(driver.provider),
+            entity_edges=edges,
+        )
 
 
 async def extract_nodes_and_edges_bulk(
     clients: GraphitiClients,
     episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]],
     edge_type_map: dict[tuple[str, str], list[str]],
-    entity_types: dict[str, BaseModel] | None = None,
+    entity_types: dict[str, type[BaseModel]] | None = None,
     excluded_entity_types: list[str] | None = None,
-    edge_types: dict[str, BaseModel] | None = None,
+    edge_types: dict[str, type[BaseModel]] | None = None,
 ) -> tuple[list[list[EntityNode]], list[list[EntityEdge]]]:
     extracted_nodes_bulk: list[list[EntityNode]] = await semaphore_gather(
         *[
@@ -204,85 +288,113 @@ async def dedupe_nodes_bulk(
     clients: GraphitiClients,
     extracted_nodes: list[list[EntityNode]],
     episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]],
-    entity_types: dict[str, BaseModel] | None = None,
+    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]
-    )
+    """Resolve entity duplicates across an in-memory batch using a two-pass strategy.
 
-    # 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
-
-                # 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
 
@@ -292,7 +404,7 @@ async def dedupe_edges_bulk(
     extracted_edges: list[list[EntityEdge]],
     episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]],
     _entities: list[EntityNode],
-    edge_types: dict[str, BaseModel],
+    edge_types: dict[str, type[BaseModel]],
     _edge_type_map: dict[tuple[str, str], list[str]],
 ) -> dict[str, list[EntityEdge]]:
     embedder = clients.embedder
@@ -307,16 +419,23 @@ async def dedupe_edges_bulk(
     dedupe_tuples: list[tuple[EpisodicNode, EntityEdge, list[EntityEdge]]] = []
     for i, edges_i in enumerate(extracted_edges):
         existing_edges: list[EntityEdge] = []
-        for j, edges_j in enumerate(extracted_edges):
-            if i == j:
-                continue
+        for edges_j in extracted_edges:
             existing_edges += edges_j
 
         for edge in edges_i:
             candidates: list[EntityEdge] = []
             for existing_edge in existing_edges:
+                # Skip self-comparison
+                if edge.uuid == existing_edge.uuid:
+                    continue
                 # 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
+                if (
+                    edge.source_node_uuid != existing_edge.source_node_uuid
+                    or edge.target_node_uuid != existing_edge.target_node_uuid
+                ):
+                    continue
+
                 edge_words = set(edge.fact.lower().split())
                 existing_edge_words = set(existing_edge.fact.lower().split())
                 has_overlap = not edge_words.isdisjoint(existing_edge_words)
@@ -339,12 +458,19 @@ async def dedupe_edges_bulk(
     ] = await semaphore_gather(
         *[
             resolve_extracted_edge(
-                clients.llm_client, edge, candidates, candidates, episode, edge_types
+                clients.llm_client,
+                edge,
+                candidates,
+                candidates,
+                episode,
+                edge_types,
+                set(edge_types),
             )
             for episode, edge, candidates in dedupe_tuples
         ]
     )
 
+    # For now we won't track edge invalidation
     duplicate_pairs: list[tuple[str, str]] = []
     for i, (_, _, duplicates) in enumerate(bulk_edge_resolutions):
         episode, edge, candidates = dedupe_tuples[i]