main.py

"""
Copyright 2024, Zep Software, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import logging
from 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 cogram.cross_encoder.client import CrossEncoderClient
from cogram.cross_encoder.openai_reranker_client import OpenAIRerankerClient
from cogram.utils.decorators import handle_multiple_group_ids
from cogram.driver.driver import GraphDriver
from cogram.driver.neo4j_driver import Neo4jDriver
from cogram.core.edges import (
    CommunityEdge,
    Edge,
    EntityEdge,
    EpisodicEdge,
    HasEpisodeEdge,
    NextEpisodeEdge,
    create_entity_edge_embeddings,
)
from cogram.embedder import EmbedderClient, OpenAIEmbedder
from cogram.core.errors import EdgeNotFoundError, NodeNotFoundError
from cogram.core.clients import GraphitiClients
from cogram.utils.helpers import (
    get_default_group_id,
    semaphore_gather,
    validate_excluded_entity_types,
    validate_group_id,
)
from cogram.llm_client import LLMClient, OpenAIClient
from cogram.namespaces import EdgeNamespace, NodeNamespace
from cogram.core.nodes import (
    CommunityNode,
    EntityNode,
    EpisodeType,
    EpisodicNode,
    Node,
    SagaNode,
    create_entity_node_embeddings,
)
from cogram.prompts.lib import prompt_library
from cogram.prompts.summarize_sagas import SagaSummary
from cogram.search.search import SearchConfig, search
from cogram.search.search_config import DEFAULT_SEARCH_LIMIT, SearchResults
from cogram.search.search_config_recipes import (
    COMBINED_HYBRID_SEARCH_CROSS_ENCODER,
    EDGE_HYBRID_SEARCH_NODE_DISTANCE,
    EDGE_HYBRID_SEARCH_RRF,
)
from cogram.search.search_filters import SearchFilters
from cogram.search.search_utils import (
    RELEVANT_SCHEMA_LIMIT,
    get_mentioned_nodes,
)
from cogram.telemetry import capture_event
from cogram.utils.tracer import Tracer, create_tracer
from cogram.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 cogram.utils.datetime_utils import utc_now
from cogram.utils.maintenance.community_operations import (
    build_communities,
    remove_communities,
    update_community,
)
from cogram.utils.maintenance.edge_operations import (
    build_episodic_edges,
    extract_edges,
    resolve_extracted_edge,
    resolve_extracted_edges,
)
from cogram.utils.maintenance.graph_data_operations import (
    EPISODE_WINDOW_LEN,
    retrieve_episodes,
)
from cogram.utils.maintenance.node_operations import (
    extract_attributes_from_nodes,
    extract_nodes,
    resolve_extracted_nodes,
)
from cogram.utils.ontology_utils.entity_types_utils import validate_entity_types
from cogram.utils.text_utils import MAX_SUMMARY_CHARS

logger = logging.getLogger(__name__)

load_dotenv()


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 Cogram:
    """Cogram: forked Graphiti with the intent layer, Engram cache, Redis active
    memory, knot synthesis, and MCP server baked in.

    The ``Graphiti`` name is preserved as an alias at module level for backward
    compatibility — everything that previously imported ``Graphiti`` keeps working.
    """

    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',
    ):
        """
        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.
        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. Overrides SEMAPHORE_LIMIT set in the environment.
            If not set, the Graphiti default is used.
        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'.

        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 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
        if llm_client:
            self.llm_client = llm_client
        else:
            self.llm_client = OpenAIClient()
        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
        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,
        )

        # 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 token 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, now: datetime) -> 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.
        now : datetime
            The current timestamp for creation.

        Returns
        -------
        SagaNode
            The existing or newly created saga node.
        """
        from cogram.utils.helpers import parse_db_date

        records, _, _ = await self.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
            """,
            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
            )

        saga = SagaNode(name=saga_name, group_id=group_id, created_at=now)
        await saga.save(self.driver)
        return saga

    async def _saga_get_previous_episode_uuid(
        self, saga_uuid: str, current_episode_uuid: str
    ) -> str | None:
        """Find the most recent episode UUID in a saga, excluding the current one."""
        if self.driver.graph_operations_interface:
            try:
                return await self.driver.graph_operations_interface.saga_get_previous_episode_uuid(
                    self.driver, saga_uuid, current_episode_uuid
                )
            except NotImplementedError:
                pass

        records, _, _ = await self.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[str] | None:
        """Retrieve episode contents for summarization, 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.

        If the saga has been summarized before (``last_summarized_at`` is set),
        only episodes added after that timestamp 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 (or all if never summarized).
        max_episodes = 200
        since = saga.last_summarized_at

        # Try IoC interface first, fall back to raw Cypher
        episode_contents = await self._saga_get_episode_contents(
            saga_id, since=since, limit=max_episodes
        )
        if episode_contents 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 > $since
                    RETURN e.content AS content
                    ORDER BY e.valid_at ASC, e.created_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)
                    RETURN e.content AS content
                    ORDER BY e.valid_at DESC, e.created_at DESC
                    LIMIT $limit
                    """,
                    saga_uuid=saga_id,
                    limit=max_episodes,
                    routing_='r',
                )
                # Reverse to chronological order for the prompt
                records = list(reversed(records))

            episode_contents = [r['content'] for r in records if r.get('content')]

        if not episode_contents:
            logger.info(f'No new episodes found for saga {saga_id}, skipping summary')
            return saga

        context = {
            'saga_name': saga.name,
            'existing_summary': saga.summary or '',
            'episodes': episode_contents,
        }

        llm_response = await self.llm_client.generate_response(
            prompt_library.summarize_sagas.summarize_saga(context),
            response_model=SagaSummary,
            prompt_name='summarize_sagas.summarize_saga',
        )

        summary = llm_response.get('summary', '')
        if len(summary) > MAX_SUMMARY_CHARS:
            summary = summary[:MAX_SUMMARY_CHARS]

        saga.summary = summary
        saga.last_summarized_at = utc_now()
        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 _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],
        custom_extraction_instructions: str | 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)
        """
        episodes = episode if isinstance(episode, list) else [episode]
        primary_episode = episodes[0]

        extracted_edges = await extract_edges(
            self.clients,
            episode,
            extracted_nodes,
            previous_episodes,
            edge_type_map,
            group_id,
            edge_types,
            custom_extraction_instructions,
        )

        edges = resolve_edge_pointers(extracted_edges, uuid_map)

        resolved_edges, invalidated_edges, new_edges = await resolve_extracted_edges(
            self.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,
    ) -> 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.
        """
        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(
            self.driver,
            episodes,
            episodic_edges,
            nodes,
            entity_edges,
            self.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):
                saga_node = await self._get_or_create_saga(saga, group_id, now)
            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
                )

            # 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(self.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(self.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(self.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,
        custom_extraction_instructions: str | None = None,
    ) -> tuple[
        dict[str, list[EntityNode]],
        dict[str, str],
        list[list[EntityEdge]],
    ]:
        """Extract nodes and edges from all episodes and deduplicate."""
        # Extract all nodes and edges for each episode
        extracted_nodes_bulk, extracted_edges_bulk = await extract_nodes_and_edges_bulk(
            self.clients,
            episode_context,
            edge_type_map=edge_type_map,
            edge_types=edge_types,
            entity_types=entity_types,
            excluded_entity_types=excluded_entity_types,
            custom_extraction_instructions=custom_extraction_instructions,
        )

        # Dedupe extracted nodes in memory
        nodes_by_episode, uuid_map = await dedupe_nodes_bulk(
            self.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],
    ) -> tuple[list[EntityNode], list[EntityEdge], list[EntityEdge], dict[str, str]]:
        """Resolve nodes and edges against the existing graph."""
        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(
                    self.clients,
                    nodes_by_episode_unique[episode.uuid],
                    episode,
                    previous_episodes,
                    entity_types,
                )
                for episode, previous_episodes in episode_context
            ]
        )

        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(
                    self.clients,
                    nodes_by_episode_unique[episode.uuid],
                    episode,
                    previous_episodes,
                    entity_types,
                )
                for episode, previous_episodes in episode_context
            ]
        )

        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(
                    self.clients,
                    edges_by_episode_unique[episode.uuid],
                    episode,
                    final_hydrated_nodes,
                    edge_types or {},
                    edge_type_map,
                )
                for episode in episodes
            ]
        )

        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)

    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)

        if group_id is None:
            # if group_id is None, use the default group id by the provider
            # and the preset database name will be used
            group_id = get_default_group_id(self.driver.provider)
        else:
            validate_group_id(group_id)
            if group_id != self.driver._database:
                # if group_id is provided, use it as the database name
                self.driver = self.driver.clone(database=group_id)
                self.clients.driver = self.driver

        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,
                    )
                    if previous_episode_uuids is None
                    else await EpisodicNode.get_by_uuids(self.driver, previous_episode_uuids)
                )

                # Get or create episode
                episode = (
                    await EpisodicNode.get_by_uuid(self.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(
                    self.clients,
                    episode,
                    previous_episodes,
                    entity_types,
                    excluded_entity_types,
                    custom_extraction_instructions,
                )

                nodes, uuid_map, _ = await resolve_extracted_nodes(
                    self.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,
                )

                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(
                    self.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,
                )

                # Update communities if requested
                communities = []
                community_edges = []
                if update_communities:
                    communities, community_edges = await semaphore_gather(
                        *[
                            update_community(self.driver, self.llm_client, self.embedder, node)
                            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')

                results = AddEpisodeResults(
                    episode=episode,
                    episodic_edges=episodic_edges,
                    nodes=hydrated_nodes,
                    edges=entity_edges,
                    communities=communities,
                    community_edges=community_edges,
                )

                # ---- COGRAM HOOK ------------------------------------------------------
                # When COGRAM_FULL_PIPELINE=true, fire the cogram post-write pipeline:
                #   intent annotation on edges, narrative on entities, profile distill,
                #   knot synthesis, Redis cache invalidation. Runs as a background task
                #   so MCP latency stays bounded.
                # When unset, behaves like stock graphiti.
                # ----------------------------------------------------------------------
                import os as _os
                if _os.environ.get('COGRAM_FULL_PIPELINE', '').lower() == 'true':
                    try:
                        import asyncio as _asyncio
                        from cogram.pipeline.post_write import cogram_post_write
                        from cogram.core.config import Settings as _Settings
                        _asyncio.create_task(
                            cogram_post_write(
                                graphiti=self,
                                episode_result=results,
                                group_id=group_id or 'default',
                                settings=_Settings.from_env(),
                            )
                        )
                    except Exception as _hook_exc:
                        logger.warning(f'cogram_post_write hook scheduling failed: {_hook_exc}')

                return results

            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,
        custom_extraction_instructions: str | None = None,
        saga: str | SagaNode | None = None,
    ) -> 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.
        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.

        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.

        Important: This method does not perform edge invalidation or date extraction steps.
        If these operations are required, use the `add_episode` method instead for each
        individual episode.
        """
        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()

                # if group_id is None, use the default group id by the provider
                if group_id is None:
                    group_id = get_default_group_id(self.driver.provider)
                else:
                    validate_group_id(group_id)
                    if group_id != self.driver._database:
                        # if group_id is provided, use it as the database name
                        self.driver = self.driver.clone(database=group_id)
                        self.clients.driver = self.driver

                # 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(self.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=self.driver,
                    episodic_nodes=episodes,
                    episodic_edges=[],
                    entity_nodes=[],
                    entity_edges=[],
                    embedder=self.embedder,
                )

                # Get previous episode context for each episode
                episode_context = await retrieve_previous_episodes_bulk(self.driver, episodes)

                # 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,
                    custom_extraction_instructions,
                )

                # 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(
                    self.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,
                )

                # Resolved pointers for episodic edges
                resolved_episodic_edges = resolve_edge_pointers(episodic_edges, final_uuid_map)

                # save data to KG
                await add_nodes_and_edges_bulk(
                    self.driver,
                    episodes,
                    resolved_episodic_edges,
                    final_hydrated_nodes,
                    resolved_edges + invalidated_edges,
                    self.embedder,
                )

                # Handle saga association if provided
                if saga is not None:
                    # Get or create saga node based on input type
                    if isinstance(saga, str):
                        saga_node = await self._get_or_create_saga(saga, group_id, now)
                    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, ''
                    )

                    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(self.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(self.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(self.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
        await remove_communities(driver)

        community_nodes, community_edges = await build_communities(
            driver, self.llm_client, group_ids
        )

        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,
    ) -> 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,
            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
    ) -> 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

        valid_edges = await EntityEdge.get_between_nodes(
            self.driver, edge.source_node_uuid, edge.target_node_uuid
        )

        related_edges = (
            await search(
                self.clients,
                edge.fact,
                group_ids=[edge.group_id],
                config=EDGE_HYBRID_SEARCH_RRF,
                search_filter=SearchFilters(edge_uuids=[edge.uuid for edge in valid_edges]),
            )
        ).edges
        existing_edges = (
            await search(
                self.clients,
                edge.fact,
                group_ids=[edge.group_id],
                config=EDGE_HYBRID_SEARCH_RRF,
                search_filter=SearchFilters(),
            )
        ).edges

        resolved_edge, invalidated_edges, _ = await resolve_extracted_edge(
            self.llm_client,
            edge,
            related_edges,
            existing_edges,
            EpisodicNode(
                name='',
                source=EpisodeType.text,
                source_description='',
                content='',
                valid_at=edge.valid_at or utc_now(),
                entity_edges=[],
                group_id=edge.group_id,
            ),
            None,
        )

        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)

        # We should only delete edges created by the episode
        edges_to_delete: list[EntityEdge] = []
        for edge in edges:
            if edge.episodes and edge.episodes[0] == episode.uuid:
                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)

        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)


# ---------------------------------------------------------------------------
# Backward compatibility: keep the original 'Graphiti' name importable
# (any code doing 'from cogram.main import Graphiti' or 'from cogram import
# Graphiti' continues to work after the class rename to 'Cogram').
# ---------------------------------------------------------------------------
Graphiti = Cogram