topoteretes--cognee/cognee/context_global_variables.py

import os
import warnings
from contextvars import ContextVar
from typing import Optional, Union
from uuid import UUID

from cognee.base_config import get_base_config
from cognee.exceptions import CogneeValidationError
from cognee.infrastructure.llm.config import LLMConfig
from cognee.infrastructure.databases.vector.embeddings.config import EmbeddingConfig
from cognee.infrastructure.databases.vector.config import (
    get_vectordb_config,
    get_vectordb_context_config,
)

from cognee.infrastructure.files.storage.config import file_storage_config
from cognee.modules.users.methods import get_user
from cognee.infrastructure.databases.graph.config import get_graph_config, get_graph_context_config
from cognee.infrastructure.databases.utils.get_or_create_dataset_database import (
    get_or_create_dataset_database,
)
from cognee.infrastructure.databases.utils.resolve_dataset_database_connection_info import (
    resolve_dataset_database_connection_info,
)

# Note: ContextVar allows us to use different graph db configurations in Cognee
#       for different async tasks, threads and processes
vector_db_config = ContextVar("vector_db_config", default=None)
graph_db_config = ContextVar("graph_db_config", default=None)
current_dataset_id = ContextVar("current_dataset_id", default=None)
# Note: same mechanism for LLM and embedding configs so that the LiteLLM client
#       and the embedding engine can use per-context (e.g. per-request) configs.
llm_config: ContextVar[Optional[LLMConfig]] = ContextVar("llm_config", default=None)
embedding_config = ContextVar("embedding_config", default=None)
session_user = ContextVar("session_user", default=None)
# Labels the pipeline stage (extraction | summarization | query) whose LLM
# config is currently active on `llm_config`, for tracing (see pipeline_stage).
current_pipeline_stage: ContextVar[Optional[str]] = ContextVar(
    "current_pipeline_stage", default=None
)


async def set_session_user_context_variable(user):
    session_user.set(user)


def multi_user_support_possible():
    graph_db_config = get_graph_config()
    vector_db_config = get_vectordb_config()

    graph_handler = graph_db_config.graph_dataset_database_handler
    vector_handler = vector_db_config.vector_dataset_database_handler
    from cognee.infrastructure.databases.dataset_database_handler import (
        supported_dataset_database_handlers,
    )

    if graph_handler not in supported_dataset_database_handlers:
        raise EnvironmentError(
            "Unsupported graph dataset to database handler configured. Cannot add support for multi-user access control mode. Please use a supported graph dataset to database handler or set the environment variables ENABLE_BACKEND_ACCESS_CONTROL to false to switch off multi-user access control mode.\n"
            f"Selected graph dataset to database handler: {graph_handler}\n"
            f"Supported dataset to database handlers: {list(supported_dataset_database_handlers.keys())}\n"
        )

    if vector_handler not in supported_dataset_database_handlers:
        raise EnvironmentError(
            "Unsupported vector dataset to database handler configured. Cannot add support for multi-user access control mode. Please use a supported vector dataset to database handler or set the environment variables ENABLE_BACKEND_ACCESS_CONTROL to false to switch off multi-user access control mode.\n"
            f"Selected vector dataset to database handler: {vector_handler}\n"
            f"Supported dataset to database handlers: {list(supported_dataset_database_handlers.keys())}\n"
        )

    if (
        supported_dataset_database_handlers[graph_handler]["handler_provider"]
        != graph_db_config.graph_database_provider
    ):
        raise EnvironmentError(
            "The selected graph dataset to database handler does not work with the configured graph database provider. Cannot add support for multi-user access control mode. Please use a supported graph dataset to database handler or set the environment variables ENABLE_BACKEND_ACCESS_CONTROL to false to switch off multi-user access control mode.\n"
            f"Selected graph database provider: {graph_db_config.graph_database_provider}\n"
            f"Selected graph dataset to database handler: {graph_handler}\n"
            f"Supported dataset to database handlers: {list(supported_dataset_database_handlers.keys())}\n"
        )

    if (
        supported_dataset_database_handlers[vector_handler]["handler_provider"]
        != vector_db_config.vector_db_provider
    ):
        raise EnvironmentError(
            "The selected vector dataset to database handler does not work with the configured vector database provider. Cannot add support for multi-user access control mode. Please use a supported vector dataset to database handler or set the environment variables ENABLE_BACKEND_ACCESS_CONTROL to false to switch off multi-user access control mode.\n"
            f"Selected vector database provider: {vector_db_config.vector_db_provider}\n"
            f"Selected vector dataset to database handler: {vector_handler}\n"
            f"Supported dataset to database handlers: {list(supported_dataset_database_handlers.keys())}\n"
        )

    return True


def backend_access_control_enabled():
    backend_access_control = os.environ.get("ENABLE_BACKEND_ACCESS_CONTROL", None)
    if backend_access_control is None:
        # If backend access control is not defined in environment variables,
        # enable it by default if graph and vector DBs can support it, otherwise disable it
        return multi_user_support_possible()
    elif backend_access_control.lower() == "true":
        # If enabled, ensure that the current graph and vector DBs can support it
        return multi_user_support_possible()
    return False


VECTOR_DBS_WITH_MULTI_USER_SUPPORT = ["lancedb", "pgvector", "falkor"]
GRAPH_DBS_WITH_MULTI_USER_SUPPORT = ["ladybug", "kuzu", "falkor", "postgres"]


class DatabaseContextManager:
    """Dual-mode helper returned by :func:`set_database_global_context_variables`.

    Supports both ``await`` (legacy) and ``async with`` (scoped) usage.

    Note: Single-use object, should not be reused across multiple calls.
    """

    __slots__ = (
        "_dataset",
        "_user_id",
        "_llm_config",
        "_embedding_config",
        "_applied",
        "_dataset_token",
        "_llm_token",
        "_embedding_token",
    )

    def __init__(
        self,
        dataset: Union[str, UUID],
        user_id: UUID,
        llm_config: Optional[LLMConfig] = None,
        embedding_config: Optional[EmbeddingConfig] = None,
    ) -> None:
        self._dataset = dataset
        self._user_id = user_id
        self._llm_config = llm_config
        self._embedding_config = embedding_config
        self._applied = False
        self._dataset_token = None
        self._llm_token = None
        self._embedding_token = None

    async def apply_database_context_variables(
        self, dataset: Union[str, UUID], user_id: UUID
    ) -> None:
        self._dataset_token = current_dataset_id.set(str(dataset) if dataset is not None else None)

        # LLM and embedding configs are an explicit, caller-provided override and
        # are intentionally applied regardless of backend access control: callers
        # may want per-context LLM/embedding configs even in single-tenant mode.
        if self._llm_config is not None:
            self._llm_token = llm_config.set(self._llm_config)
        if self._embedding_config is not None:
            self._embedding_token = embedding_config.set(self._embedding_config)

        if not backend_access_control_enabled():
            return

        # In multi-user mode a dataset is required to resolve the per-dataset
        # database; fail fast with a clear message instead of a downstream
        # "user None" lookup error.
        if dataset is None:
            raise CogneeValidationError(
                "A dataset must be provided when backend access control is enabled."
            )

        # Imported lazily to avoid circular imports at module load.
        from cognee.infrastructure.databases.dataset_queue import dataset_queue

        await dataset_queue().ensure_slot(dataset)

        user = await get_user(user_id)

        # To ensure permissions are enforced properly all datasets will have their own databases
        dataset_database = await get_or_create_dataset_database(dataset, user)
        # Ensure that all connection info is resolved properly
        dataset_database = await resolve_dataset_database_connection_info(dataset_database)

        base_config = get_base_config()
        data_root_directory = os.path.join(
            base_config.data_root_directory, str(user.tenant_id or user.id)
        )
        databases_directory_path = os.path.join(
            base_config.system_root_directory, "databases", str(user.id)
        )

        # Set vector and graph database configuration based on dataset database information
        vector_config = {
            "vector_db_provider": dataset_database.vector_database_provider,
            "vector_db_url": dataset_database.vector_database_url,
            "vector_db_key": dataset_database.vector_database_key,
            "vector_db_name": dataset_database.vector_database_name,
            "vector_db_port": dataset_database.vector_database_connection_info.get("port", ""),
            "vector_db_host": dataset_database.vector_database_connection_info.get("host", ""),
            "vector_db_username": dataset_database.vector_database_connection_info.get(
                "username", ""
            ),
            "vector_db_password": dataset_database.vector_database_connection_info.get(
                "password", ""
            ),
            # Inherit subprocess mode from the global config so that per-dataset DB wrappers
            # are also spawned as subprocesses when the feature is enabled.
            "vector_db_subprocess_enabled": get_vectordb_config().vector_db_subprocess_enabled,
        }

        graph_config = {
            "graph_database_provider": dataset_database.graph_database_provider,
            "graph_database_url": dataset_database.graph_database_url,
            "graph_database_name": dataset_database.graph_database_name,
            "graph_database_key": dataset_database.graph_database_key,
            "graph_file_path": os.path.join(
                databases_directory_path, dataset_database.graph_database_name
            ),
            "graph_database_username": dataset_database.graph_database_connection_info.get(
                "graph_database_username", ""
            ),
            "graph_database_password": dataset_database.graph_database_connection_info.get(
                "graph_database_password", ""
            ),
            "graph_database_host": dataset_database.graph_database_connection_info.get(
                "graph_database_host", ""
            ),
            "graph_database_allow_anonymous": dataset_database.graph_database_connection_info.get(
                "graph_database_allow_anonymous",
                get_graph_config().graph_database_allow_anonymous,
            ),
            "graph_dataset_database_handler": dataset_database.graph_dataset_database_handler,
            "graph_database_port": dataset_database.graph_database_connection_info.get(
                "graph_database_port", ""
            ),
            # Inherit subprocess mode and Kuzu tuning from the global config so that
            # per-dataset DB wrappers are spawned with matching settings.
            "graph_database_subprocess_enabled": get_graph_config().graph_database_subprocess_enabled,
            "kuzu_num_threads": get_graph_config().kuzu_num_threads,
            "kuzu_buffer_pool_size": get_graph_config().kuzu_buffer_pool_size,
            "kuzu_max_db_size": get_graph_config().kuzu_max_db_size,
        }

        storage_config = {
            "data_root_directory": data_root_directory,
        }

        # Use ContextVar to use these graph and vector configurations are used
        # in the current async context across Cognee. Unlike the LLM/embedding
        # overrides these intentionally persist after async-with exit: callers
        # read the per-dataset databases right after a pipeline run, outside
        # this context manager.
        graph_db_config.set(graph_config)
        vector_db_config.set(vector_config)
        file_storage_config.set(storage_config)

    async def _apply(self) -> None:
        if self._applied:
            return
        await self.apply_database_context_variables(self._dataset, self._user_id)
        self._applied = True

    def __await__(self):
        # Legacy ``await set_database_global_context_variables(...)`` call shape.
        # Deprecated — emit a warning pointing users at the ``async with`` form,
        # which scopes the dataset queue slot to a block and releases it on exit.
        warnings.warn(
            "`await set_database_global_context_variables(...)` is deprecated. "
            "Use `async with set_database_global_context_variables(...):` instead "
            "for scoped dataset-queue slot management; the `async with` form "
            "releases the database queue slot automatically on block exit, the deprecated method does not.",
            DeprecationWarning,
            stacklevel=2,
        )
        return self._apply().__await__()

    async def __aenter__(self) -> "DatabaseContextManager":
        await self._apply()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        # Restore the caller-provided LLM/embedding overrides and the dataset id
        # so they don't leak into the surrounding async context. The dataset
        # graph/vector/file-storage configs are left in place on purpose (see
        # apply_database_context_variables).
        for context_var, token_attr in (
            (embedding_config, "_embedding_token"),
            (llm_config, "_llm_token"),
            (current_dataset_id, "_dataset_token"),
        ):
            token = getattr(self, token_attr)
            if token is not None:
                context_var.reset(token)
                setattr(self, token_attr, None)

        if not backend_access_control_enabled():
            return None

        from cognee.infrastructure.databases.dataset_queue import dataset_queue

        await dataset_queue().release_slot_for(self._dataset)


def set_database_global_context_variables(
    dataset: Union[str, UUID],
    user_id: UUID,
    llm_config: Optional[LLMConfig] = None,
    embedding_config: Optional[EmbeddingConfig] = None,
) -> "DatabaseContextManager":
    """Returns a dual-mode helper that is both awaitable and an async context manager.

    - ``await set_database_global_context_variables(ds, user_id)`` — legacy;
      applies the context and relies on task-end queue cleanup to release the
      slot.
    - ``async with set_database_global_context_variables(ds, user_id):`` —
      applies the context on enter; explicitly releases this dataset's queue
      slot on exit. Preferred for sequential multi-dataset loops and scoped
      operations.

    If backend access control is enabled this ensures all datasets have their
    own databases, access to which will be enforced by given permissions.
    Database name will be determined by dataset and the appropriate vector and
    graph database handlers will be enforced.

    Additionally, this acts as the queue gate for dataset-level operations:
    applying the context ensures the current asyncio task holds a
    :class:`DatasetQueue` slot for ``dataset``. Repeated calls in the same
    task for the same dataset are no-ops;. The dataset queue slot is released automatically when the
    task completes (legacy mode) or on async-with exit (scoped mode).

    If ``llm_config`` and/or ``embedding_config`` are provided they are set on
    their respective ContextVars and picked up by ``get_llm_client`` (LiteLLM)
    and ``get_embedding_engine`` in the current async context. Unlike the
    graph/vector configs these are applied even when backend access control is
    disabled, since they are an explicit caller-provided override. In the
    ``async with`` form they are restored to their prior values on exit, while
    the per-dataset graph/vector/file-storage configs persist after exit so
    callers can keep reading the dataset databases.

    Args:
        dataset: Cognee dataset name or id
        user_id: UUID of the owner of the dataset
        llm_config: Optional ``LLMConfig`` to use for LLM calls in this context.
        embedding_config: Optional ``EmbeddingConfig`` to use for embedding calls
            in this context.

    Returns:
        A :class:`DatabaseContextManager` that can be awaited or used as an
        async context manager.
    """
    return DatabaseContextManager(dataset, user_id, llm_config, embedding_config)