""" 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 json import logging from datetime import datetime from time import time from uuid import uuid4 from dotenv import load_dotenv from pydantic import BaseModel from typing_extensions import LiteralString from graphiti_core.cross_encoder.client import CrossEncoderClient from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient from graphiti_core.decorators import handle_multiple_group_ids from graphiti_core.driver.driver import GraphDriver from graphiti_core.driver.neo4j_driver import Neo4jDriver from graphiti_core.edges import ( CommunityEdge, Edge, EntityEdge, EpisodicEdge, HasEpisodeEdge, NextEpisodeEdge, create_entity_edge_embeddings, ) from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder from graphiti_core.errors import EdgeNotFoundError, NodeNotFoundError from graphiti_core.graphiti_types import ( GraphitiClients, generate_prompt_response, resolve_prompt_messages, ) from graphiti_core.helpers import ( get_default_group_id, semaphore_gather, validate_excluded_entity_types, validate_group_id, ) from graphiti_core.llm_client import LLMClient, LLMRuntime, OpenAIClient from graphiti_core.namespaces import EdgeNamespace, NodeNamespace from graphiti_core.nodes import ( CommunityNode, EntityNode, EpisodeType, EpisodicNode, Node, SagaNode, create_entity_node_embeddings, ) from graphiti_core.prompts import prompt_library from graphiti_core.prompts.lib import ( ChatPromptLibrary, ensure_prompt_library_wrapped, validate_prompt_library, ) from graphiti_core.prompts.summarize_sagas import SagaSummary from graphiti_core.search.search import SearchConfig, search from graphiti_core.search.search_config import DEFAULT_SEARCH_LIMIT, SearchResults from graphiti_core.search.search_config_recipes import ( COMBINED_HYBRID_SEARCH_CROSS_ENCODER, EDGE_HYBRID_SEARCH_NODE_DISTANCE, EDGE_HYBRID_SEARCH_RRF, ) from graphiti_core.search.search_filters import SearchFilters from graphiti_core.search.search_utils import ( RELEVANT_SCHEMA_LIMIT, get_mentioned_nodes, ) from graphiti_core.telemetry import capture_event from graphiti_core.tracer import Tracer, create_tracer from graphiti_core.utils.bulk_utils import ( RawEpisode, add_nodes_and_edges_bulk, dedupe_edges_bulk, dedupe_nodes_bulk, extract_nodes_and_edges_bulk, resolve_edge_pointers, retrieve_previous_episodes_bulk, ) from graphiti_core.utils.datetime_utils import utc_now from graphiti_core.utils.maintenance.community_operations import ( build_communities, remove_communities, update_community, ) from graphiti_core.utils.maintenance.dedup_helpers import _normalize_string_exact from graphiti_core.utils.maintenance.edge_operations import ( _pair_search_filter, _same_pair_edges, build_episodic_edges, extract_edges, resolve_extracted_edge, resolve_extracted_edges, ) from graphiti_core.utils.maintenance.graph_data_operations import ( EPISODE_WINDOW_LEN, retrieve_episodes, ) from graphiti_core.utils.maintenance.node_operations import ( _extract_entity_summaries_batch, extract_attributes_from_nodes, extract_nodes, resolve_extracted_nodes, ) from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types from graphiti_core.utils.text_utils import ( SAGA_SUMMARY_MAX_CHARS, truncate_at_sentence, truncate_utf8_to_bytes, ) logger = logging.getLogger(__name__) load_dotenv() # Episode text is packed under 100k characters. The assembled prompt adds the # existing summary (up to SAGA_SUMMARY_MAX_CHARS), template, and schema on top # of that, then shrinks until the full prompt is under 128 KiB. Request 8192 # output tokens. Persist at most SAGA_SUMMARY_MAX_CHARS. SAGA_SUMMARY_MAX_EPISODE_CHARS = 100_000 SAGA_SUMMARY_MAX_EPISODE_BYTES = SAGA_SUMMARY_MAX_EPISODE_CHARS SAGA_SUMMARY_MAX_PROMPT_BYTES = 128 * 1024 MAX_SAGA_EPISODES_FOR_SUMMARY = 20 SUMMARY_REBUILD_BATCH_SIZE = 10 SAGA_SUMMARY_MAX_TOKENS = 8192 _SAGA_SUMMARY_SCHEMA_WRAPPER = '\n\nRespond with a JSON object in the following format:\n\n' def _saga_summary_messages_utf8_bytes(messages: list) -> int: return sum(len(message.content.encode('utf-8')) for message in messages) def select_saga_summary_episodes( episodes_data: list[tuple[str, datetime | None, datetime]], max_episode_chars: int = SAGA_SUMMARY_MAX_EPISODE_CHARS, ) -> tuple[list[tuple[str, datetime | None, datetime]], int, bool]: """Select a contiguous ingestion-time episode batch using a character budget. Packs oldest-first so last_summarized_at can advance through a backlog. When the oldest episode alone exceeds the budget, truncate it rather than skip, matching Go entity-summary behavior for a single oversized episode. """ selected: list[tuple[str, datetime | None, datetime]] = [] selected_chars = 0 dropped = False for content, valid_at, created_at in episodes_data: content_chars = len(content) if selected_chars + content_chars <= max_episode_chars: selected.append((content, valid_at, created_at)) selected_chars += content_chars continue if not selected: truncated = content[:max_episode_chars] if truncated: selected.append((truncated, valid_at, created_at)) selected_chars = len(truncated) dropped = True break # Preserve chronological coherence for incremental summaries. Once the # next episode does not fit, later episodes are deferred to a future run. dropped = True break selected_bytes = sum(len(content.encode('utf-8')) for content, _, _ in selected) return selected, selected_bytes, dropped class AddEpisodeResults(BaseModel): episode: EpisodicNode episodic_edges: list[EpisodicEdge] nodes: list[EntityNode] edges: list[EntityEdge] communities: list[CommunityNode] community_edges: list[CommunityEdge] class AddBulkEpisodeResults(BaseModel): episodes: list[EpisodicNode] episodic_edges: list[EpisodicEdge] nodes: list[EntityNode] edges: list[EntityEdge] communities: list[CommunityNode] community_edges: list[CommunityEdge] class AddTripletResults(BaseModel): nodes: list[EntityNode] edges: list[EntityEdge] class Graphiti: def __init__( self, uri: str | None = None, user: str | None = None, password: str | None = None, llm_client: LLMClient | None = None, embedder: EmbedderClient | None = None, cross_encoder: CrossEncoderClient | None = None, store_raw_episode_content: bool = True, graph_driver: GraphDriver | None = None, max_coroutines: int | None = None, tracer: Tracer | None = None, trace_span_prefix: str = 'graphiti', prompt_library: ChatPromptLibrary | None = None, llm_runtime: LLMRuntime | None = None, ): """ Initialize a Graphiti instance. This constructor sets up a connection to a graph database and initializes the LLM client for natural language processing tasks. Parameters ---------- uri : str The URI of the Neo4j database. user : str The username for authenticating with the Neo4j database. password : str The password for authenticating with the Neo4j database. llm_client : LLMClient | None, optional An instance of LLMClient for natural language processing tasks. If not provided, a default OpenAIClient will be initialized. Cannot be combined with ``llm_runtime``. embedder : EmbedderClient | None, optional An instance of EmbedderClient for embedding tasks. If not provided, a default OpenAIEmbedder will be initialized. cross_encoder : CrossEncoderClient | None, optional An instance of CrossEncoderClient for reranking tasks. If not provided, a default OpenAIRerankerClient will be initialized. store_raw_episode_content : bool, optional Whether to store the raw content of episodes. Defaults to True. graph_driver : GraphDriver | None, optional An instance of GraphDriver for database operations. If not provided, a default Neo4jDriver will be initialized. max_coroutines : int | None, optional The maximum number of concurrent operations allowed at each `semaphore_gather` call site across the Graphiti ingestion and search paths. When set, this value is propagated to all internal `semaphore_gather` calls via `GraphitiClients.max_coroutines`, making this the recommended way to tune concurrency on a per-instance basis. If not set, the value falls back to the `SEMAPHORE_LIMIT` environment variable (if defined) or the built-in default (`DEFAULT_SEMAPHORE_LIMIT`). Note: this is a per-`semaphore_gather` cap, not a global cap on in-flight coroutines. Each `semaphore_gather` call allocates its own `asyncio.Semaphore`, so when a parent gather awaits child coroutines that themselves contain a gather, the inner gathers each bound at `max_coroutines` independently. For the search path that means the inner per-scope dispatches (`edge_search`, `node_search`, `community_search`) can together admit up to ~`scopes × max_coroutines` in-flight tasks, where `scopes` is the number of search scopes configured. Tune `max_coroutines` with that in mind when sizing for tight LLM rate-limits. Carve-out: when a `GraphDriver` implements `graph_operations_interface.get_community_clusters`, that path is serial inside the driver and does not honour `max_coroutines`. Community-build concurrency outside the cluster lookup (LLM summarization fan-out) is bounded as documented above. tracer : Tracer | None, optional An OpenTelemetry tracer instance for distributed tracing. If not provided, tracing is disabled (no-op). trace_span_prefix : str, optional Prefix to prepend to all span names. Defaults to 'graphiti'. prompt_library : ChatPromptLibrary | None, optional An instance-scoped prompt library used for all LLM prompt construction. If not provided, Graphiti uses its legacy module-level prompt functions. ``self.prompt_library`` and ``self.clients.prompt_library`` remain ``None`` by default. For partial customization, compose overrides with ``create_prompt_library(overrides)`` and pass the result here. Override callables must return ``ChatPrompt``. Cannot be combined with ``llm_runtime``. llm_runtime : LLMRuntime | None, optional An opt-in runtime that routes prompts across one or more ``LLMTransport`` instances. It requires a default ``LLMModel``. The runtime owns prompt selection and model routing. It has no ``small_id`` field. ``model_size`` does not select a smaller model on the runtime path. Route lightweight prompts to a smaller model. GLiNER2 does not support per-prompt routing. Cannot be combined with ``llm_client`` or ``prompt_library``. Returns ------- None Notes ----- This method establishes a connection to a graph database (Neo4j by default) using the provided credentials. It also sets up the LLM client, either using the provided client or by creating a default OpenAIClient. The default database name is defined during the driver’s construction. If a different database name is required, it should be specified in the URI or set separately after initialization. The OpenAI API key is expected to be set in the environment variables. Make sure to set the OPENAI_API_KEY environment variable before initializing Graphiti if you're using the default OpenAIClient. """ if llm_runtime is not None: extras = [ name for name, value in (('llm_client', llm_client), ('prompt_library', prompt_library)) if value is not None ] if extras: raise ValueError('llm_runtime cannot be combined with: ' + ', '.join(extras) + '.') if graph_driver: self.driver = graph_driver else: if uri is None: raise ValueError('uri must be provided when graph_driver is None') self.driver = Neo4jDriver(uri, user, password) self.store_raw_episode_content = store_raw_episode_content self.max_coroutines = max_coroutines self.llm_runtime = llm_runtime if llm_runtime is not None: self.llm_client = llm_runtime.client self.prompt_library = llm_runtime.library else: if llm_client: self.llm_client = llm_client else: self.llm_client = OpenAIClient() if prompt_library is not None: validate_prompt_library(prompt_library) self.prompt_library = ensure_prompt_library_wrapped(prompt_library) else: self.prompt_library = None if embedder: self.embedder = embedder else: self.embedder = OpenAIEmbedder() if cross_encoder: self.cross_encoder = cross_encoder else: self.cross_encoder = OpenAIRerankerClient() # Initialize tracer self.tracer = create_tracer(tracer, trace_span_prefix) # Set tracer on clients if self.llm_runtime is not None: self.llm_runtime.set_tracer(self.tracer) else: self.llm_client.set_tracer(self.tracer) self.clients = GraphitiClients( driver=self.driver, llm_client=self.llm_client, embedder=self.embedder, cross_encoder=self.cross_encoder, tracer=self.tracer, prompt_library=self.prompt_library, llm_runtime=self.llm_runtime, max_coroutines=self.max_coroutines, ) # Initialize namespace API (graphiti.nodes.entity.save(), etc.) self.nodes = NodeNamespace(self.driver, self.embedder) self.edges = EdgeNamespace(self.driver, self.embedder) # Capture telemetry event self._capture_initialization_telemetry() def _capture_initialization_telemetry(self): """Capture telemetry event for Graphiti initialization.""" try: # Detect provider types from class names llm_provider = self._get_provider_type(self.llm_client) embedder_provider = self._get_provider_type(self.embedder) reranker_provider = self._get_provider_type(self.cross_encoder) database_provider = self._get_provider_type(self.driver) properties = { 'llm_provider': llm_provider, 'embedder_provider': embedder_provider, 'reranker_provider': reranker_provider, 'database_provider': database_provider, } capture_event('graphiti_initialized', properties) except Exception: # Silently handle telemetry errors pass @property def token_tracker(self): """Access the LLM client's token usage tracker. Returns the TokenUsageTracker from the LLM client, which can be used to: - Get token usage by prompt type: tracker.get_usage() - Get total usage: tracker.get_total_usage() - Print a formatted summary: tracker.print_summary() - Reset tracking: tracker.reset() """ return self.llm_client.token_tracker def _get_provider_type(self, client) -> str: """Get provider type from client class name.""" if client is None: return 'none' class_name = client.__class__.__name__.lower() # LLM providers if 'openai' in class_name: return 'openai' elif 'azure' in class_name: return 'azure' elif 'anthropic' in class_name: return 'anthropic' elif 'crossencoder' in class_name: return 'crossencoder' elif 'gemini' in class_name: return 'gemini' elif 'groq' in class_name: return 'groq' # Database providers elif 'neo4j' in class_name: return 'neo4j' elif 'falkor' in class_name: return 'falkordb' # Embedder providers elif 'voyage' in class_name: return 'voyage' else: return 'unknown' async def close(self): """ Close the connection to the Neo4j database. This method safely closes the driver connection to the Neo4j database. It should be called when the Graphiti instance is no longer needed or when the application is shutting down. Parameters ---------- self Returns ------- None Notes ----- It's important to close the driver connection to release system resources and ensure that all pending transactions are completed or rolled back. This method should be called as part of a cleanup process, potentially in a context manager or a shutdown hook. Example: graphiti = Graphiti(uri, user, password) try: # Use graphiti... finally: graphiti.close() """ await self.driver.close() async def _get_or_create_saga( self, saga_name: str, group_id: str, created_at: datetime, driver: GraphDriver | None = None, ) -> SagaNode: """ Get an existing saga by name or create a new one. Parameters ---------- saga_name : str The name of the saga. group_id : str The group id for the saga. created_at : datetime Timestamp to stamp on a newly created saga. Callers should pass the originating episode's reference time (``valid_at``) rather than the current wall-clock time so the saga's ``created_at`` reflects the episode it was minted from. Returns ------- SagaNode The existing or newly created saga node. """ from graphiti_core.helpers import parse_db_date driver = driver or self.driver records, _, _ = await driver.execute_query( """ MATCH (s:Saga {name: $name, group_id: $group_id}) RETURN s.uuid AS uuid, s.name AS name, s.group_id AS group_id, s.created_at AS created_at, s.first_episode_uuid AS first_episode_uuid, s.last_episode_uuid AS last_episode_uuid """, name=saga_name, group_id=group_id, routing_='r', ) if records: record = records[0] return SagaNode( uuid=record['uuid'], name=record['name'], group_id=record['group_id'], created_at=parse_db_date(record['created_at']), # type: ignore first_episode_uuid=record['first_episode_uuid'], last_episode_uuid=record['last_episode_uuid'], ) saga = SagaNode(name=saga_name, group_id=group_id, created_at=created_at) await saga.save(driver) return saga async def _saga_get_previous_episode_uuid( self, saga_uuid: str, current_episode_uuid: str, driver: GraphDriver | None = None ) -> str | None: """Find the most recent episode UUID in a saga, excluding the current one.""" driver = driver or self.driver if driver.graph_operations_interface: try: return await driver.graph_operations_interface.saga_get_previous_episode_uuid( driver, saga_uuid, current_episode_uuid ) except NotImplementedError: pass records, _, _ = await driver.execute_query( """ MATCH (s:Saga {uuid: $saga_uuid})-[:HAS_EPISODE]->(e:Episodic) WHERE e.uuid <> $current_episode_uuid RETURN e.uuid AS uuid ORDER BY e.valid_at DESC, e.created_at DESC LIMIT 1 """, saga_uuid=saga_uuid, current_episode_uuid=current_episode_uuid, routing_='r', ) if records: return records[0]['uuid'] return None async def _saga_get_episode_contents( self, saga_uuid: str, since: datetime | None = None, limit: int = 200, ) -> list[tuple[str, datetime | None, datetime]] | None: """Retrieve (content, valid_at, created_at) episodes, using IoC if available.""" if self.driver.graph_operations_interface: try: return await self.driver.graph_operations_interface.saga_get_episode_contents( self.driver, saga_uuid, since=since, limit=limit ) except NotImplementedError: pass return None async def summarize_saga( self, saga_id: str, ) -> SagaNode: """Incrementally summarize a saga using only new episodes since the last summary. Two watermarks are maintained on the saga node, with deliberately different semantics: - ``last_summarized_at`` is the ingestion-time *filter* watermark: the next run picks up any episode whose ``created_at`` is greater than this value. It advances only to the latest episode actually included in the summary prompt, so byte-budgeted episodes remain reachable on subsequent runs. - ``last_summarized_episode_valid_at`` is the *temporal* watermark: the maximum ``valid_at`` (episode reference time) across the episodes covered by the current summary. Consumers asking "how recent is the content of this summary in event-time?" should use this field, not ``last_summarized_at``. If the saga has been summarized before, only episodes added after ``last_summarized_at`` are fetched. The existing summary is provided to the LLM as context so no information is lost. On the first call (no prior summary), all episodes are included. Parameters ---------- saga_id : str The UUID of the saga to summarize. Returns ------- SagaNode The updated saga node with the new summary. Raises ------ NodeNotFoundError If the saga with the given UUID does not exist. """ saga = await SagaNode.get_by_uuid(self.driver, saga_id) # Fetch only episodes added since the last summary. Cap the window to the same episode # count node summaries use so a long thread or document cannot fill # the model context window in one pass. Deferred episodes remain # eligible because last_summarized_at advances only through this batch. max_episodes = MAX_SAGA_EPISODES_FOR_SUMMARY since = saga.last_summarized_at existing_summary = truncate_at_sentence(saga.summary or '', SAGA_SUMMARY_MAX_CHARS) # Try IoC interface first, fall back to raw Cypher episodes_data = await self._saga_get_episode_contents( saga_id, since=since, limit=max_episodes ) if episodes_data is None: if since is not None: records, _, _ = await self.driver.execute_query( """ MATCH (s:Saga {uuid: $saga_uuid})-[:HAS_EPISODE]->(e:Episodic) WHERE e.created_at IS NOT NULL AND e.created_at > $since RETURN e.content AS content, e.valid_at AS valid_at, e.created_at AS created_at ORDER BY e.created_at ASC, e.valid_at ASC LIMIT $limit """, saga_uuid=saga_id, since=since, limit=max_episodes, routing_='r', ) else: records, _, _ = await self.driver.execute_query( """ MATCH (s:Saga {uuid: $saga_uuid})-[:HAS_EPISODE]->(e:Episodic) WHERE e.created_at IS NOT NULL RETURN e.content AS content, e.valid_at AS valid_at, e.created_at AS created_at ORDER BY e.created_at ASC, e.valid_at ASC LIMIT $limit """, saga_uuid=saga_id, limit=max_episodes, routing_='r', ) from graphiti_core.helpers import parse_db_date episodes_data = [] for r in records: content = r.get('content') created_at = parse_db_date(r.get('created_at')) if content and created_at is not None: episodes_data.append((content, parse_db_date(r.get('valid_at')), created_at)) if not episodes_data: saga._summary_episodes_fetched = 0 saga._summary_episodes_selected = 0 saga._summary_episodes_skipped = 0 saga._summary_selected_episode_bytes = 0 logger.info(f'No new episodes found for saga {saga_id}, skipping summary') return saga total_episode_count = len(episodes_data) episodes_data, selected_episode_bytes, episodes_dropped = select_saga_summary_episodes( episodes_data, max_episode_chars=SAGA_SUMMARY_MAX_EPISODE_CHARS ) saga._summary_episodes_fetched = total_episode_count saga._summary_episodes_selected = len(episodes_data) saga._summary_episodes_skipped = total_episode_count - len(episodes_data) saga._summary_selected_episode_bytes = selected_episode_bytes if not episodes_data: logger.info(f'No saga episode content fit the summary byte budget for saga {saga_id}') return saga if episodes_dropped: logger.info( 'Saga summary episode payload constrained by byte budget', extra={ 'saga_id': saga_id, 'episodes_fetched': total_episode_count, 'episodes_selected': len(episodes_data), 'selected_episode_bytes': selected_episode_bytes, 'max_episode_chars': SAGA_SUMMARY_MAX_EPISODE_CHARS, 'max_prompt_bytes': SAGA_SUMMARY_MAX_PROMPT_BYTES, 'episodes_dropped': episodes_dropped, }, ) def _prompt_context( selected: list[tuple[str, datetime | None, datetime]], ) -> dict[str, object]: prompt_episodes_data = sorted( selected, key=lambda episode: (episode[1] or episode[2], episode[2]) ) return { 'saga_name': saga.name, 'existing_summary': existing_summary, 'episodes': [content for content, _, _ in prompt_episodes_data], } schema_reserve = len(_SAGA_SUMMARY_SCHEMA_WRAPPER.encode('utf-8')) + len( json.dumps(SagaSummary.model_json_schema()).encode('utf-8') ) saga_clients = getattr(self, 'clients', None) while episodes_data: context = _prompt_context(episodes_data) prompt_bytes = ( _saga_summary_messages_utf8_bytes( resolve_prompt_messages( 'summarize_sagas.summarize_saga', prompt_library.summarize_sagas.summarize_saga, context, clients=saga_clients, ) ) + schema_reserve ) if prompt_bytes <= SAGA_SUMMARY_MAX_PROMPT_BYTES: break if len(episodes_data) == 1: content, valid_at, created_at = episodes_data[0] over = prompt_bytes - SAGA_SUMMARY_MAX_PROMPT_BYTES truncated = truncate_utf8_to_bytes( content, max(0, len(content.encode('utf-8')) - over) ) episodes_data = [(truncated, valid_at, created_at)] if truncated else [] break episodes_data = episodes_data[:-1] episodes_dropped = True saga._summary_episodes_selected = len(episodes_data) saga._summary_episodes_skipped = total_episode_count - len(episodes_data) saga._summary_selected_episode_bytes = sum( len(content.encode('utf-8')) for content, _, _ in episodes_data ) if not episodes_data: logger.info(f'No saga episode content fit the summary byte budget for saga {saga_id}') return saga valid_ats = [valid_at for _, valid_at, _ in episodes_data if valid_at is not None] summarized_created_ats = [created_at for _, _, created_at in episodes_data] context = _prompt_context(episodes_data) llm_response = await generate_prompt_response( self.llm_client, 'summarize_sagas.summarize_saga', prompt_library.summarize_sagas.summarize_saga, context, clients=saga_clients, response_model=SagaSummary, max_tokens=SAGA_SUMMARY_MAX_TOKENS, ) summary = llm_response.get('summary', '') if len(summary) > SAGA_SUMMARY_MAX_CHARS: summary = truncate_at_sentence(summary, SAGA_SUMMARY_MAX_CHARS) saga.summary = summary # Ingestion-time watermark for the next-run filter. Advance only # through the episodes that were actually sent to the summary prompt so # byte-budgeted later episodes remain reachable on subsequent runs. saga.last_summarized_at = max(summarized_created_ats) if summarized_created_ats else None # Episode-time watermark for public/temporal consumers: advance only # forward to the latest reference time we just summarized. If no # episode in this batch carried a valid_at, leave the previous value unchanged. if valid_ats: new_episode_watermark = max(valid_ats) if ( saga.last_summarized_episode_valid_at is None or new_episode_watermark > saga.last_summarized_episode_valid_at ): saga.last_summarized_episode_valid_at = new_episode_watermark await saga.save(self.driver) logger.info(f'Updated summary for saga {saga_id}') return saga async def build_indices_and_constraints(self, delete_existing: bool = False): """ Build indices and constraints in the Neo4j database. This method sets up the necessary indices and constraints in the Neo4j database to optimize query performance and ensure data integrity for the knowledge graph. Parameters ---------- self delete_existing : bool, optional Whether to clear existing indices before creating new ones. Returns ------- None Notes ----- This method should typically be called once during the initial setup of the knowledge graph or when updating the database schema. It uses the driver's `build_indices_and_constraints` method to perform the actual database operations. The specific indices and constraints created depend on the implementation of the driver's `build_indices_and_constraints` method. Refer to the specific driver documentation for details on the exact database schema modifications. Caution: Running this method on a large existing database may take some time and could impact database performance during execution. """ await self.driver.build_indices_and_constraints(delete_existing) async def _extract_and_resolve_nodes( self, episode: EpisodicNode | list[EpisodicNode], previous_episodes: list[EpisodicNode], entity_types: dict[str, type[BaseModel]] | None, excluded_entity_types: list[str] | None, ) -> tuple[ list[EntityNode], dict[str, str], list[tuple[EntityNode, EntityNode]], dict[str, list[int]] ]: """Extract nodes from episode(s) and resolve against existing graph.""" episodes = episode if isinstance(episode, list) else [episode] primary_episode = episodes[0] extracted_nodes, node_episode_index_map = await extract_nodes( self.clients, episode, previous_episodes, entity_types, excluded_entity_types ) nodes, uuid_map, duplicates = await resolve_extracted_nodes( self.clients, extracted_nodes, primary_episode, previous_episodes, entity_types, ) return nodes, uuid_map, duplicates, node_episode_index_map async def _merge_materialized_endpoint_nodes( self, materialized_nodes: list[EntityNode], episode: EpisodicNode, previous_episodes: list[EpisodicNode], entity_types: dict[str, type[BaseModel]] | None, nodes: list[EntityNode], uuid_map: dict[str, str], clients: GraphitiClients | None = None, ) -> None: """Fold endpoints materialized by ``extract_edges`` into the resolved node set. These names never went through node extraction, so they still have to be deduplicated the way ordinary extracted nodes are. Nodes resolved for this episode are not persisted yet and are therefore invisible to graph search, so match against them by name first and send only the remainder through ``resolve_extracted_nodes``. ``nodes`` and ``uuid_map`` are updated in place. """ clients = clients or self.clients nodes_by_name = {_normalize_string_exact(node.name): node for node in nodes} known_uuids = {node.uuid for node in nodes} unresolved: list[EntityNode] = [] for materialized_node in materialized_nodes: if materialized_node.uuid in known_uuids: continue existing = nodes_by_name.get(_normalize_string_exact(materialized_node.name)) if existing is not None: uuid_map[materialized_node.uuid] = existing.uuid continue unresolved.append(materialized_node) if not unresolved: return resolved_nodes, materialized_uuid_map, _ = await resolve_extracted_nodes( clients, unresolved, episode, previous_episodes, entity_types, ) uuid_map.update(materialized_uuid_map) for resolved_node in resolved_nodes: if resolved_node.uuid in known_uuids: continue nodes.append(resolved_node) known_uuids.add(resolved_node.uuid) async def _extract_and_resolve_edges( self, episode: EpisodicNode | list[EpisodicNode], extracted_nodes: list[EntityNode], previous_episodes: list[EpisodicNode], edge_type_map: dict[tuple[str, str], list[str]], group_id: str, edge_types: dict[str, type[BaseModel]] | None, nodes: list[EntityNode], uuid_map: dict[str, str], strict_edge_types: bool = False, custom_extraction_instructions: str | None = None, entity_types: dict[str, type[BaseModel]] | None = None, clients: GraphitiClients | None = None, ) -> tuple[list[EntityEdge], list[EntityEdge], list[EntityEdge]]: """Extract edges from episode(s) and resolve against existing graph. Returns ------- tuple[list[EntityEdge], list[EntityEdge], list[EntityEdge]] A tuple of (resolved_edges, invalidated_edges, new_edges) where: - resolved_edges: All edges after resolution - invalidated_edges: Edges invalidated by new information - new_edges: Only edges that are new to the graph (not duplicates) """ clients = clients or self.clients episodes = episode if isinstance(episode, list) else [episode] primary_episode = episodes[0] extracted_edges, materialized_nodes = await extract_edges( clients, episode, extracted_nodes, previous_episodes, edge_type_map, group_id, edge_types, strict_edge_types, custom_extraction_instructions, ) # Merge only nodes that extract_edges newly materialized — never the full # pre-resolution extracted_nodes list, which also contains deduped duplicates # whose uuid_map remaps must be preserved for resolve_edge_pointers. if materialized_nodes: await self._merge_materialized_endpoint_nodes( materialized_nodes, primary_episode, previous_episodes, entity_types, nodes, uuid_map, clients=clients, ) edges = resolve_edge_pointers(extracted_edges, uuid_map) resolved_edges, invalidated_edges, new_edges = await resolve_extracted_edges( clients, edges, primary_episode, nodes, edge_types or {}, edge_type_map, ) return resolved_edges, invalidated_edges, new_edges async def _process_episode_data( self, episode: EpisodicNode | list[EpisodicNode], nodes: list[EntityNode], entity_edges: list[EntityEdge], now: datetime, group_id: str, saga: str | SagaNode | None = None, saga_previous_episode_uuid: str | None = None, node_episode_index_map: dict[str, list[int]] | None = None, clients: GraphitiClients | None = None, ) -> tuple[list[EpisodicEdge], EpisodicNode]: """Process and save episode data to the graph. Parameters ---------- episode : EpisodicNode | list[EpisodicNode] The episode(s) to process. nodes : list[EntityNode] The entity nodes extracted from the episode(s). entity_edges : list[EntityEdge] The entity edges extracted from the episode(s). now : datetime The current timestamp. group_id : str The group id for the episode. saga : str | SagaNode | None Optional. Either a saga name (str) or a SagaNode object to associate this episode with. If a string is provided, the saga will be looked up by name or created if it doesn't exist. saga_previous_episode_uuid : str | None Optional. UUID of the previous episode in the saga. If provided, skips the database query to find the most recent episode. Useful for efficiently adding multiple episodes to the same saga in sequence. node_episode_index_map : dict[str, list[int]] | None Optional mapping from node UUID to 0-indexed episode positions for building episodic edges with correct attribution. clients : GraphitiClients | None Optional request-scoped clients bundle. Defaults to ``self.clients``. Callers pass a per-request bundle so concurrent calls for different group_ids target the correct database (issue #1676). """ clients = clients or self.clients driver = clients.driver episodes = episode if isinstance(episode, list) else [episode] episode_uuids = [ep.uuid for ep in episodes] episodic_edges = build_episodic_edges(nodes, episode_uuids, now, node_episode_index_map) for ep in episodes: ep.entity_edges = [edge.uuid for edge in entity_edges] if not self.store_raw_episode_content: ep.content = '' await add_nodes_and_edges_bulk( driver, episodes, episodic_edges, nodes, entity_edges, clients.embedder, ) primary_episode = episodes[0] # Handle saga association if provided if saga is not None: # Get or create saga node based on input type if isinstance(saga, str): # Use the originating episode's reference time (valid_at) for a # newly created saga so its created_at matches the episode that # minted it, not the wall-clock time of this run. saga_created_at = primary_episode.valid_at or now saga_node = await self._get_or_create_saga( saga, group_id, saga_created_at, driver=driver ) else: saga_node = saga # Use provided previous episode UUID or query for it previous_episode_uuid: str | None = saga_previous_episode_uuid if previous_episode_uuid is None: previous_episode_uuid = await self._saga_get_previous_episode_uuid( saga_node.uuid, primary_episode.uuid, driver=driver ) # Create NEXT_EPISODE edge from the previous episode to the new one if previous_episode_uuid is not None: next_episode_edge = NextEpisodeEdge( source_node_uuid=previous_episode_uuid, target_node_uuid=primary_episode.uuid, group_id=group_id, created_at=now, ) await next_episode_edge.save(driver) # Create HAS_EPISODE edge from saga to the new episode has_episode_edge = HasEpisodeEdge( source_node_uuid=saga_node.uuid, target_node_uuid=primary_episode.uuid, group_id=group_id, created_at=now, ) await has_episode_edge.save(driver) # Track first and last episode on the saga node if saga_node.first_episode_uuid is None: saga_node.first_episode_uuid = primary_episode.uuid saga_node.last_episode_uuid = primary_episode.uuid await saga_node.save(driver) return episodic_edges, primary_episode async def _extract_and_dedupe_nodes_bulk( self, episode_context: list[tuple[EpisodicNode, list[EpisodicNode]]], edge_type_map: dict[tuple[str, str], list[str]], edge_types: dict[str, type[BaseModel]] | None, entity_types: dict[str, type[BaseModel]] | None, excluded_entity_types: list[str] | None, strict_edge_types: bool = False, custom_extraction_instructions: str | None = None, use_combined_extraction: bool = True, clients: GraphitiClients | None = None, ) -> tuple[ dict[str, list[EntityNode]], dict[str, str], list[list[EntityEdge]], ]: """Extract nodes and edges from all episodes and deduplicate.""" clients = clients or self.clients # Extract all nodes and edges for each episode extracted_nodes_bulk, extracted_edges_bulk = await extract_nodes_and_edges_bulk( clients, episode_context, edge_type_map=edge_type_map, edge_types=edge_types, entity_types=entity_types, excluded_entity_types=excluded_entity_types, strict_edge_types=strict_edge_types, custom_extraction_instructions=custom_extraction_instructions, use_combined_extraction=use_combined_extraction, ) # Dedupe extracted nodes in memory nodes_by_episode, uuid_map = await dedupe_nodes_bulk( clients, extracted_nodes_bulk, episode_context, entity_types ) return nodes_by_episode, uuid_map, extracted_edges_bulk async def _resolve_nodes_and_edges_bulk( self, nodes_by_episode: dict[str, list[EntityNode]], edges_by_episode: dict[str, list[EntityEdge]], episode_context: list[tuple[EpisodicNode, list[EpisodicNode]]], entity_types: dict[str, type[BaseModel]] | None, edge_types: dict[str, type[BaseModel]] | None, edge_type_map: dict[tuple[str, str], list[str]], episodes: list[EpisodicNode], clients: GraphitiClients | None = None, ) -> tuple[list[EntityNode], list[EntityEdge], list[EntityEdge], dict[str, str]]: """Resolve nodes and edges against the existing graph.""" clients = clients or self.clients nodes_by_uuid: dict[str, EntityNode] = { node.uuid: node for nodes in nodes_by_episode.values() for node in nodes } # Get unique nodes per episode nodes_by_episode_unique: dict[str, list[EntityNode]] = {} nodes_uuid_set: set[str] = set() for episode, _ in episode_context: nodes_by_episode_unique[episode.uuid] = [] nodes = [nodes_by_uuid[node.uuid] for node in nodes_by_episode[episode.uuid]] for node in nodes: if node.uuid not in nodes_uuid_set: nodes_by_episode_unique[episode.uuid].append(node) nodes_uuid_set.add(node.uuid) # Resolve nodes node_results = await semaphore_gather( *[ resolve_extracted_nodes( clients, nodes_by_episode_unique[episode.uuid], episode, previous_episodes, entity_types, ) for episode, previous_episodes in episode_context ], max_coroutines=self.max_coroutines, ) resolved_nodes: list[EntityNode] = [] uuid_map: dict[str, str] = {} for result in node_results: resolved_nodes.extend(result[0]) uuid_map.update(result[1]) # Update nodes_by_uuid with resolved nodes for resolved_node in resolved_nodes: nodes_by_uuid[resolved_node.uuid] = resolved_node # Update nodes_by_episode_unique with resolved pointers for episode_uuid, nodes in nodes_by_episode_unique.items(): updated_nodes: list[EntityNode] = [] for node in nodes: updated_node_uuid = uuid_map.get(node.uuid, node.uuid) updated_node = nodes_by_uuid[updated_node_uuid] updated_nodes.append(updated_node) nodes_by_episode_unique[episode_uuid] = updated_nodes # Extract attributes for resolved nodes hydrated_nodes_results: list[list[EntityNode]] = await semaphore_gather( *[ extract_attributes_from_nodes( clients, nodes_by_episode_unique[episode.uuid], episode, previous_episodes, entity_types, ) for episode, previous_episodes in episode_context ], max_coroutines=self.max_coroutines, ) final_hydrated_nodes = [node for nodes in hydrated_nodes_results for node in nodes] # Resolve edges with updated pointers edges_by_episode_unique: dict[str, list[EntityEdge]] = {} edges_uuid_set: set[str] = set() for episode_uuid, edges in edges_by_episode.items(): edges_with_updated_pointers = resolve_edge_pointers(edges, uuid_map) edges_by_episode_unique[episode_uuid] = [] for edge in edges_with_updated_pointers: if edge.uuid not in edges_uuid_set: edges_by_episode_unique[episode_uuid].append(edge) edges_uuid_set.add(edge.uuid) edge_results = await semaphore_gather( *[ resolve_extracted_edges( clients, edges_by_episode_unique[episode.uuid], episode, final_hydrated_nodes, edge_types or {}, edge_type_map, ) for episode in episodes ], max_coroutines=self.max_coroutines, ) resolved_edges: list[EntityEdge] = [] invalidated_edges: list[EntityEdge] = [] for result in edge_results: resolved_edges.extend(result[0]) invalidated_edges.extend(result[1]) # result[2] is new_edges - not used in bulk flow since attributes # are extracted before edge resolution return final_hydrated_nodes, resolved_edges, invalidated_edges, uuid_map @handle_multiple_group_ids async def retrieve_episodes( self, reference_time: datetime, last_n: int = EPISODE_WINDOW_LEN, group_ids: list[str] | None = None, source: EpisodeType | None = None, driver: GraphDriver | None = None, saga: str | None = None, ) -> list[EpisodicNode]: """ Retrieve the last n episodic nodes from the graph. This method fetches a specified number of the most recent episodic nodes from the graph, relative to the given reference time. Parameters ---------- reference_time : datetime The reference time to retrieve episodes before. last_n : int, optional The number of episodes to retrieve. Defaults to EPISODE_WINDOW_LEN. group_ids : list[str | None], optional The group ids to return data from. source : EpisodeType | None, optional Filter episodes by source type. driver : GraphDriver | None, optional The graph driver to use. If not provided, uses the default driver. saga : str | None, optional If provided, only retrieve episodes that belong to the saga with this name. Returns ------- list[EpisodicNode] A list of the most recent EpisodicNode objects. Notes ----- The actual retrieval is performed by the `retrieve_episodes` function from the `graphiti_core.utils` module, unless a saga is specified. """ if driver is None: driver = self.clients.driver if driver.graph_operations_interface: try: return await driver.graph_operations_interface.retrieve_episodes( driver, reference_time, last_n, group_ids, source, saga ) except NotImplementedError: pass return await retrieve_episodes(driver, reference_time, last_n, group_ids, source, saga) def _resolve_request_scope( self, group_id: str | None ) -> tuple[str, GraphDriver, GraphitiClients]: """Resolve the effective ``group_id`` and a request-scoped driver/clients bundle. Historically ``add_episode`` / ``add_episode_bulk`` reassigned the shared ``self.driver`` (and ``self.clients.driver``) when a ``group_id`` mapped to a different database. Because those coroutines contain many ``await`` points (LLM calls, embeddings, DB writes), a concurrent call for a *different* ``group_id`` could reassign ``self.driver`` mid-execution, causing the first call's remaining operations to target the wrong database and silently persist episodes under the wrong graph (issue #1676). Instead of mutating shared instance state, this returns a per-call driver and a matching ``GraphitiClients`` copy that callers thread through the request. The default-database case reuses the shared instances unchanged (no copy, no clone), preserving existing behaviour for the common single-database deployment. """ if group_id is None: # Use the provider's default group id and the preset database name. return get_default_group_id(self.driver.provider), self.driver, self.clients validate_group_id(group_id) if group_id == self.driver._database: return group_id, self.driver, self.clients driver = self.driver.clone(database=group_id) clients = self.clients.model_copy(update={'driver': driver}) return group_id, driver, clients async def add_episode( self, name: str, episode_body: str, source_description: str, reference_time: datetime, source: EpisodeType = EpisodeType.message, group_id: str | None = None, uuid: str | None = None, update_communities: bool = False, entity_types: dict[str, type[BaseModel]] | None = None, excluded_entity_types: list[str] | None = None, previous_episode_uuids: list[str] | None = None, edge_types: dict[str, type[BaseModel]] | None = None, edge_type_map: dict[tuple[str, str], list[str]] | None = None, custom_extraction_instructions: str | None = None, saga: str | SagaNode | None = None, saga_previous_episode_uuid: str | None = None, ) -> AddEpisodeResults: """ Process an episode and update the graph. This method extracts information from the episode, creates nodes and edges, and updates the graph database accordingly. Parameters ---------- name : str The name of the episode. episode_body : str The content of the episode. source_description : str A description of the episode's source. reference_time : datetime The reference time for the episode. source : EpisodeType, optional The type of the episode. Defaults to EpisodeType.message. group_id : str | None An id for the graph partition the episode is a part of. uuid : str | None Optional uuid of the episode. update_communities : bool Optional. Whether to update communities with new node information entity_types : dict[str, BaseModel] | None Optional. Dictionary mapping entity type names to their Pydantic model definitions. excluded_entity_types : list[str] | None Optional. List of entity type names to exclude from the graph. Entities classified into these types will not be added to the graph. Can include 'Entity' to exclude the default entity type. previous_episode_uuids : list[str] | None Optional. list of episode uuids to use as the previous episodes. If this is not provided, the most recent episodes by created_at date will be used. custom_extraction_instructions : str | None Optional. Custom extraction instructions string to be included in the extract entities and extract edges prompts. This allows for additional instructions or context to guide the extraction process. saga : str | SagaNode | None Optional. Either a saga name (str) or a SagaNode object to associate this episode with. If a string is provided and a saga with this name already exists in the group, the episode will be added to it. Otherwise, a new saga will be created. Sagas are connected to episodes via HAS_EPISODE edges, and consecutive episodes are linked via NEXT_EPISODE edges. saga_previous_episode_uuid : str | None Optional. UUID of the previous episode in the saga. If provided, skips the database query to find the most recent episode. Useful for efficiently adding multiple episodes to the same saga in sequence. The returned AddEpisodeResults.episode.uuid can be passed as this parameter for the next episode. Returns ------- None Notes ----- This method performs several steps including node extraction, edge extraction, deduplication, and database updates. It also handles embedding generation and edge invalidation. It is recommended to run this method as a background process, such as in a queue. It's important that each episode is added sequentially and awaited before adding the next one. For web applications, consider using FastAPI's background tasks or a dedicated task queue like Celery for this purpose. Example using FastAPI background tasks: @app.post("/add_episode") async def add_episode_endpoint(episode_data: EpisodeData): background_tasks.add_task(graphiti.add_episode, **episode_data.dict()) return {"message": "Episode processing started"} """ start = time() now = utc_now() validate_entity_types(entity_types) validate_excluded_entity_types(excluded_entity_types, entity_types) # Resolve a request-scoped driver/clients bundle rather than mutating the # shared ``self.driver``, so concurrent calls for different group_ids cannot # clobber each other's database target mid-execution (issue #1676). group_id, driver, clients = self._resolve_request_scope(group_id) with self.tracer.start_span('add_episode') as span: try: # Retrieve previous episodes for context previous_episodes = ( await self.retrieve_episodes( reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id], source=source, driver=driver, ) if previous_episode_uuids is None else await EpisodicNode.get_by_uuids(driver, previous_episode_uuids) ) # Get or create episode episode = ( await EpisodicNode.get_by_uuid(driver, uuid) if uuid is not None else EpisodicNode( name=name, group_id=group_id, labels=[], source=source, content=episode_body, source_description=source_description, created_at=now, valid_at=reference_time, ) ) # Create default edge type map edge_type_map_default = ( {('Entity', 'Entity'): list(edge_types.keys())} if edge_types is not None else {('Entity', 'Entity'): []} ) # Extract and resolve nodes extracted_nodes, node_episode_index_map = await extract_nodes( clients, episode, previous_episodes, entity_types, excluded_entity_types, custom_extraction_instructions, ) nodes, uuid_map, _ = await resolve_extracted_nodes( clients, extracted_nodes, episode, previous_episodes, entity_types, ) # Extract and resolve edges in parallel with attribute extraction ( resolved_edges, invalidated_edges, new_edges, ) = await self._extract_and_resolve_edges( episode, extracted_nodes, previous_episodes, edge_type_map or edge_type_map_default, group_id, edge_types, nodes, uuid_map, custom_extraction_instructions=custom_extraction_instructions, entity_types=entity_types, clients=clients, ) entity_edges = resolved_edges + invalidated_edges # Extract node attributes - only pass new edges for summary generation # to avoid duplicating facts that already exist in the graph hydrated_nodes = await extract_attributes_from_nodes( clients, nodes, episode, previous_episodes, entity_types, edges=new_edges, ) # Process and save episode data (including saga association if provided) episodic_edges, episode = await self._process_episode_data( episode, hydrated_nodes, entity_edges, now, group_id, saga, saga_previous_episode_uuid, node_episode_index_map, clients=clients, ) # Update communities if requested communities = [] community_edges = [] if update_communities: communities, community_edges = await semaphore_gather( *[ update_community( driver, clients.llm_client, clients.embedder, node, clients=clients, ) for node in nodes ], max_coroutines=self.max_coroutines, ) end = time() # Add span attributes span.add_attributes( { 'episode.uuid': episode.uuid, 'episode.source': source.value, 'episode.reference_time': reference_time.isoformat(), 'group_id': group_id, 'node.count': len(hydrated_nodes), 'edge.count': len(entity_edges), 'edge.invalidated_count': len(invalidated_edges), 'previous_episodes.count': len(previous_episodes), 'entity_types.count': len(entity_types) if entity_types else 0, 'edge_types.count': len(edge_types) if edge_types else 0, 'update_communities': update_communities, 'communities.count': len(communities) if update_communities else 0, 'duration_ms': (end - start) * 1000, } ) logger.info(f'Completed add_episode in {(end - start) * 1000} ms') return AddEpisodeResults( episode=episode, episodic_edges=episodic_edges, nodes=hydrated_nodes, edges=entity_edges, communities=communities, community_edges=community_edges, ) except Exception as e: span.set_status('error', str(e)) span.record_exception(e) raise e async def add_episode_bulk( self, bulk_episodes: list[RawEpisode], group_id: str | None = None, entity_types: dict[str, type[BaseModel]] | None = None, excluded_entity_types: list[str] | None = None, edge_types: dict[str, type[BaseModel]] | None = None, edge_type_map: dict[tuple[str, str], list[str]] | None = None, strict_edge_types: bool = False, custom_extraction_instructions: str | None = None, saga: str | SagaNode | None = None, use_combined_extraction: bool = True, ) -> AddBulkEpisodeResults: """ Process multiple episodes in bulk and update the graph. This method extracts information from multiple episodes, creates nodes and edges, and updates the graph database accordingly, all in a single batch operation. Parameters ---------- bulk_episodes : list[RawEpisode] A list of RawEpisode objects to be processed and added to the graph. group_id : str | None An id for the graph partition the episode is a part of. entity_types : dict[str, type[BaseModel]] | None Optional. A dictionary mapping entity type names to Pydantic models. excluded_entity_types : list[str] | None Optional. A list of entity type names to exclude from extraction. edge_types : dict[str, type[BaseModel]] | None Optional. A dictionary mapping edge type names to Pydantic models. edge_type_map : dict[tuple[str, str], list[str]] | None Optional. A mapping of (source_type, target_type) to allowed edge types. strict_edge_types : bool Optional. Whether to discard extracted edges that do not use one of the provided edge_types. custom_extraction_instructions : str | None Optional. Custom extraction instructions string to be included in the extract entities and extract edges prompts. This allows for additional instructions or context to guide the extraction process. saga : str | SagaNode | None Optional. Either a saga name (str) or a SagaNode object to associate all episodes with. If a string is provided and a saga with this name already exists in the group, the episodes will be added to it. Otherwise, a new saga will be created. Sagas are connected to episodes via HAS_EPISODE edges, and consecutive episodes are linked via NEXT_EPISODE edges. use_combined_extraction : bool If True, use one combined node-and-edge LLM call per episode. If False, use separate node and edge extraction. Returns ------- AddBulkEpisodeResults Notes ----- This method performs several steps including: - Saving all episodes to the database - Retrieving previous episode context for each new episode - Extracting nodes and edges from all episodes - Generating embeddings for nodes and edges - Deduplicating nodes and edges - Saving nodes, episodic edges, and entity edges to the knowledge graph This bulk operation is designed for efficiency when processing multiple episodes at once. However, it's important to ensure that the bulk operation doesn't overwhelm system resources. Consider implementing rate limiting or chunking for very large batches of episodes. Edge invalidation and date extraction (``valid_at`` / ``invalid_at``) are performed in the bulk path as well: edges flow through ``extract_edges`` and ``resolve_extracted_edges`` just like in ``add_episode``, and any invalidated edges are persisted alongside the newly resolved ones. """ with self.tracer.start_span('add_episode_bulk') as bulk_span: bulk_span.add_attributes({'episode.count': len(bulk_episodes)}) try: start = time() now = utc_now() # Resolve a request-scoped driver/clients bundle rather than mutating # the shared ``self.driver``, so concurrent calls for different # group_ids cannot clobber each other's database target (issue #1676). group_id, driver, clients = self._resolve_request_scope(group_id) # Create default edge type map edge_type_map_default = ( {('Entity', 'Entity'): list(edge_types.keys())} if edge_types is not None else {('Entity', 'Entity'): []} ) episodes = [ await EpisodicNode.get_by_uuid(driver, episode.uuid) if episode.uuid is not None else EpisodicNode( name=episode.name, labels=[], source=episode.source, content=episode.content, source_description=episode.source_description, group_id=group_id, created_at=now, valid_at=episode.reference_time, ) for episode in bulk_episodes ] # Save all episodes await add_nodes_and_edges_bulk( driver=driver, episodic_nodes=episodes, episodic_edges=[], entity_nodes=[], entity_edges=[], embedder=clients.embedder, ) # Get previous episode context for each episode episode_context = await retrieve_previous_episodes_bulk( driver, episodes, max_coroutines=self.max_coroutines ) # Extract and dedupe nodes and edges ( nodes_by_episode, uuid_map, extracted_edges_bulk, ) = await self._extract_and_dedupe_nodes_bulk( episode_context, edge_type_map or edge_type_map_default, edge_types, entity_types, excluded_entity_types, strict_edge_types=strict_edge_types, custom_extraction_instructions=custom_extraction_instructions, use_combined_extraction=use_combined_extraction, clients=clients, ) # Create Episodic Edges episodic_edges: list[EpisodicEdge] = [] for episode_uuid, nodes in nodes_by_episode.items(): episodic_edges.extend(build_episodic_edges(nodes, episode_uuid, now)) # Re-map edge pointers and dedupe edges extracted_edges_bulk_updated: list[list[EntityEdge]] = [ resolve_edge_pointers(edges, uuid_map) for edges in extracted_edges_bulk ] edges_by_episode = await dedupe_edges_bulk( clients, extracted_edges_bulk_updated, episode_context, [], edge_types or {}, edge_type_map or edge_type_map_default, ) # Resolve nodes and edges against the existing graph ( final_hydrated_nodes, resolved_edges, invalidated_edges, final_uuid_map, ) = await self._resolve_nodes_and_edges_bulk( nodes_by_episode, edges_by_episode, episode_context, entity_types, edge_types, edge_type_map or edge_type_map_default, episodes, clients=clients, ) # Resolved pointers for episodic edges resolved_episodic_edges = resolve_edge_pointers(episodic_edges, final_uuid_map) edges_to_save = resolved_edges + invalidated_edges for episode in episodes: episode.entity_edges = list( dict.fromkeys( edge.uuid for edge in edges_to_save if episode.uuid in edge.episodes ) ) # save data to KG await add_nodes_and_edges_bulk( driver, episodes, resolved_episodic_edges, final_hydrated_nodes, edges_to_save, clients.embedder, ) # Handle saga association if provided if saga is not None: # Get or create saga node based on input type if isinstance(saga, str): # Anchor a newly minted saga to the earliest episode # reference time in the bulk, so created_at reflects the # episode window rather than the time this run started. valid_ats = [ep.valid_at for ep in episodes if ep.valid_at is not None] saga_created_at = min(valid_ats) if valid_ats else now saga_node = await self._get_or_create_saga( saga, group_id, saga_created_at, driver=driver ) else: saga_node = saga # Sort episodes by valid_at to create NEXT_EPISODE chain in correct order sorted_episodes = sorted(episodes, key=lambda e: e.valid_at) # Find the most recent episode already in the saga previous_episode_uuid = await self._saga_get_previous_episode_uuid( saga_node.uuid, '', driver=driver ) for episode in sorted_episodes: # Create NEXT_EPISODE edge from the previous episode if previous_episode_uuid is not None: next_episode_edge = NextEpisodeEdge( source_node_uuid=previous_episode_uuid, target_node_uuid=episode.uuid, group_id=group_id, created_at=now, ) await next_episode_edge.save(driver) # Create HAS_EPISODE edge from saga to episode has_episode_edge = HasEpisodeEdge( source_node_uuid=saga_node.uuid, target_node_uuid=episode.uuid, group_id=group_id, created_at=now, ) await has_episode_edge.save(driver) # Update previous_episode_uuid for the next iteration previous_episode_uuid = episode.uuid # Track first and last episode on the saga node if sorted_episodes: if saga_node.first_episode_uuid is None: saga_node.first_episode_uuid = sorted_episodes[0].uuid saga_node.last_episode_uuid = sorted_episodes[-1].uuid await saga_node.save(driver) end = time() # Add span attributes bulk_span.add_attributes( { 'group_id': group_id, 'node.count': len(final_hydrated_nodes), 'edge.count': len(resolved_edges + invalidated_edges), 'duration_ms': (end - start) * 1000, } ) logger.info(f'Completed add_episode_bulk in {(end - start) * 1000} ms') return AddBulkEpisodeResults( episodes=episodes, episodic_edges=resolved_episodic_edges, nodes=final_hydrated_nodes, edges=resolved_edges + invalidated_edges, communities=[], community_edges=[], ) except Exception as e: bulk_span.set_status('error', str(e)) bulk_span.record_exception(e) raise e @handle_multiple_group_ids async def build_communities( self, group_ids: list[str] | None = None, driver: GraphDriver | None = None ) -> tuple[list[CommunityNode], list[CommunityEdge]]: """ Use a community clustering algorithm to find communities of nodes. Create community nodes summarising the content of these communities. ---------- group_ids : list[str] | None Optional. Create communities only for the listed group_ids. If blank the entire graph will be used. """ if driver is None: driver = self.clients.driver # Clear existing communities (scoped: only the groups being rebuilt) await remove_communities(driver, group_ids=group_ids) # Prefer the resolved clients bundle (llm_runtime / prompt_library), # but honor an explicit driver override when provided. clients = self.clients if driver is not None and driver is not self.driver: clients = self.clients.model_copy(update={'driver': driver}) community_nodes, community_edges = await build_communities( clients.driver, clients.llm_client, group_ids, clients=clients, max_coroutines=self.max_coroutines, ) await semaphore_gather( *[node.generate_name_embedding(self.embedder) for node in community_nodes], max_coroutines=self.max_coroutines, ) await semaphore_gather( *[node.save(driver) for node in community_nodes], max_coroutines=self.max_coroutines, ) await semaphore_gather( *[edge.save(driver) for edge in community_edges], max_coroutines=self.max_coroutines, ) return community_nodes, community_edges @handle_multiple_group_ids async def search( self, query: str, center_node_uuid: str | None = None, group_ids: list[str] | None = None, num_results=DEFAULT_SEARCH_LIMIT, search_filter: SearchFilters | None = None, driver: GraphDriver | None = None, ) -> list[EntityEdge]: """ Perform a hybrid search on the knowledge graph. This method executes a search query on the graph, combining vector and text-based search techniques to retrieve relevant facts, returning the edges as a string. This is our basic out-of-the-box search, for more robust results we recommend using our more advanced search method graphiti.search_(). Parameters ---------- query : str The search query string. center_node_uuid: str, optional Facts will be reranked based on proximity to this node group_ids : list[str | None] | None, optional The graph partitions to return data from. num_results : int, optional The maximum number of results to return. Defaults to 10. Returns ------- list A list of EntityEdge objects that are relevant to the search query. Notes ----- This method uses a SearchConfig with num_episodes set to 0 and num_results set to the provided num_results parameter. The search is performed using the current date and time as the reference point for temporal relevance. """ search_config = ( EDGE_HYBRID_SEARCH_RRF if center_node_uuid is None else EDGE_HYBRID_SEARCH_NODE_DISTANCE ) search_config.limit = num_results edges = ( await search( self.clients, query, group_ids, search_config, search_filter if search_filter is not None else SearchFilters(), driver=driver, center_node_uuid=center_node_uuid, ) ).edges return edges async def _search( self, query: str, config: SearchConfig, group_ids: list[str] | None = None, center_node_uuid: str | None = None, bfs_origin_node_uuids: list[str] | None = None, search_filter: SearchFilters | None = None, ) -> SearchResults: """DEPRECATED""" return await self.search_( query, config, group_ids, center_node_uuid, bfs_origin_node_uuids, search_filter ) @handle_multiple_group_ids async def search_( self, query: str, config: SearchConfig = COMBINED_HYBRID_SEARCH_CROSS_ENCODER, group_ids: list[str] | None = None, center_node_uuid: str | None = None, bfs_origin_node_uuids: list[str] | None = None, search_filter: SearchFilters | None = None, driver: GraphDriver | None = None, query_vector: list[float] | None = None, ) -> SearchResults: """search_ (replaces _search) is our advanced search method that returns Graph objects (nodes and edges) rather than a list of facts. This endpoint allows the end user to utilize more advanced features such as filters and different search and reranker methodologies across different layers in the graph. For different config recipes refer to search/search_config_recipes. """ return await search( self.clients, query, group_ids, config, search_filter if search_filter is not None else SearchFilters(), center_node_uuid, bfs_origin_node_uuids, query_vector=query_vector, driver=driver, ) async def get_nodes_and_edges_by_episode(self, episode_uuids: list[str]) -> SearchResults: episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids) edges_list = await semaphore_gather( *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes], max_coroutines=self.max_coroutines, ) edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst] nodes = await get_mentioned_nodes(self.driver, episodes) return SearchResults(edges=edges, nodes=nodes) async def add_triplet( self, source_node: EntityNode, edge: EntityEdge, target_node: EntityNode, episode: EpisodicNode | None = None, ) -> AddTripletResults: if source_node.name_embedding is None: await source_node.generate_name_embedding(self.embedder) if target_node.name_embedding is None: await target_node.generate_name_embedding(self.embedder) if edge.fact_embedding is None: await edge.generate_embedding(self.embedder) try: resolved_source = await EntityNode.get_by_uuid(self.driver, source_node.uuid) except NodeNotFoundError: resolved_source_nodes, _, _ = await resolve_extracted_nodes( self.clients, [source_node], ) resolved_source = resolved_source_nodes[0] try: resolved_target = await EntityNode.get_by_uuid(self.driver, target_node.uuid) except NodeNotFoundError: resolved_target_nodes, _, _ = await resolve_extracted_nodes( self.clients, [target_node], ) resolved_target = resolved_target_nodes[0] nodes = [resolved_source, resolved_target] # Merge user-provided properties from original nodes into resolved nodes (excluding uuid) # Update attributes dictionary (merge rather than replace) if source_node.attributes: resolved_source.attributes.update(source_node.attributes) if target_node.attributes: resolved_target.attributes.update(target_node.attributes) # Update summary if provided by user (non-empty string) if source_node.summary: resolved_source.summary = source_node.summary if target_node.summary: resolved_target.summary = target_node.summary # Update labels (merge with existing) if source_node.labels: resolved_source.labels = list(set(resolved_source.labels) | set(source_node.labels)) if target_node.labels: resolved_target.labels = list(set(resolved_target.labels) | set(target_node.labels)) edge.source_node_uuid = resolved_source.uuid edge.target_node_uuid = resolved_target.uuid # Check if an edge with this UUID already exists with different source/target nodes. # If so, generate a new UUID to create a new edge instead of overwriting. try: existing_edge = await EntityEdge.get_by_uuid(self.driver, edge.uuid) # Edge exists - check if source/target nodes match if ( existing_edge.source_node_uuid != edge.source_node_uuid or existing_edge.target_node_uuid != edge.target_node_uuid ): # Source/target mismatch - generate new UUID to create a new edge old_uuid = edge.uuid edge.uuid = str(uuid4()) logger.info( f'Edge UUID {old_uuid} already exists with different source/target nodes. ' f'Generated new UUID {edge.uuid} to avoid overwriting.' ) except EdgeNotFoundError: # Edge doesn't exist yet, proceed normally pass # Duplicate candidates: a hybrid search pre-filtered to edges between # the pair's endpoints (either orientation) -- the backend ranks over # the pair's edges only, never a graph-wide scan. The post-filter # drops self-loops the {a,b}x{a,b} filter admits. Mirrors # resolve_extracted_edges. related_edges = _same_pair_edges( edge, ( await search( self.clients, edge.fact, group_ids=[edge.group_id], config=EDGE_HYBRID_SEARCH_RRF, search_filter=_pair_search_filter(edge), query_vector=edge.fact_embedding, ) ).edges, ) existing_edges = ( await search( self.clients, edge.fact, group_ids=[edge.group_id], config=EDGE_HYBRID_SEARCH_RRF, search_filter=SearchFilters(), query_vector=edge.fact_embedding, ) ).edges # Keep edges that are already duplicate candidates out of the # invalidation list so the LLM is not shown the same fact under two # indices. Mirrors resolve_extracted_edges' overlap dedup. related_uuids = {related_edge.uuid for related_edge in related_edges} existing_edges = [e for e in existing_edges if e.uuid not in related_uuids] dedupe_episode = ( episode if episode is not None else EpisodicNode( name='', source=EpisodeType.text, source_description='', content='', valid_at=edge.valid_at or utc_now(), entity_edges=[], group_id=edge.group_id, ) ) resolved_edge, invalidated_edges, _ = await resolve_extracted_edge( self.llm_client, edge, related_edges, existing_edges, dedupe_episode, None, clients=self.clients, ) edges: list[EntityEdge] = [resolved_edge] + invalidated_edges await create_entity_edge_embeddings(self.embedder, edges) await create_entity_node_embeddings(self.embedder, nodes) await add_nodes_and_edges_bulk(self.driver, [], [], nodes, edges, self.embedder) return AddTripletResults(edges=edges, nodes=nodes) async def remove_episode(self, episode_uuid: str): # Find the episode to be deleted episode = await EpisodicNode.get_by_uuid(self.driver, episode_uuid) # Find edges mentioned by the episode edges = await EntityEdge.get_by_uuids(self.driver, episode.entity_edges) # Keep edges that other episodes also support. edges_to_delete: list[EntityEdge] = [] edges_to_update: list[EntityEdge] = [] for edge in edges: if episode.uuid in edge.episodes: edge.episodes = [uuid for uuid in edge.episodes if uuid != episode.uuid] if edge.episodes: edges_to_update.append(edge) else: edges_to_delete.append(edge) # Find nodes mentioned by the episode nodes = await get_mentioned_nodes(self.driver, [episode]) # We should delete all nodes that are only mentioned in the deleted episode nodes_to_delete: list[EntityNode] = [] for node in nodes: query: LiteralString = 'MATCH (e:Episodic)-[:MENTIONS]->(n:Entity {uuid: $uuid}) RETURN count(*) AS episode_count' records, _, _ = await self.driver.execute_query(query, uuid=node.uuid, routing_='r') for record in records: if record['episode_count'] == 1: nodes_to_delete.append(node) deleted_node_uuids = {node.uuid for node in nodes_to_delete} surviving_nodes = [node for node in nodes if node.uuid not in deleted_node_uuids] for edge in edges_to_update: await edge.load_fact_embedding(self.driver) await edge.save(self.driver) await Edge.delete_by_uuids(self.driver, [edge.uuid for edge in edges_to_delete]) await Node.delete_by_uuids(self.driver, [node.uuid for node in nodes_to_delete]) await episode.delete(self.driver) async def rebuild_node_summary(node: EntityNode) -> None: old_summary = node.summary try: episodes = await EpisodicNode.get_by_entity_node_uuid(self.driver, node.uuid) remaining_episodes = sorted( [remaining for remaining in episodes if remaining.uuid != episode.uuid], key=lambda remaining: remaining.valid_at or remaining.created_at, ) node.summary = '' for batch_start in range(0, len(remaining_episodes), SUMMARY_REBUILD_BATCH_SIZE): episode_batch = remaining_episodes[ batch_start : batch_start + SUMMARY_REBUILD_BATCH_SIZE ] await _extract_entity_summaries_batch( self.llm_client, [node], episode_batch, None, None, {}, skip_fact_appending=True, clients=self.clients, ) await node.load_name_embedding(self.driver) await node.save(self.driver) except Exception: node.summary = old_summary logger.exception( 'Failed to rebuild the summary of entity %s after removal of episode %s', node.uuid, episode.uuid, ) await semaphore_gather( *[rebuild_node_summary(node) for node in surviving_nodes], max_coroutines=self.max_coroutines, )