app/cognee_runtime/bootstrap.py

"""Cognee runtime bootstrap for ContextFirewall.

Single source of truth for configuring the Cognee SDK. Config is read from the
environment (loaded from ``backend/.env`` in dev; real env vars on the Hugging
Face Space in prod) and applied via ``cognee.config`` setters BEFORE any
add/cognify/search/memify/forget call.

Profiles (selected purely by which env vars are present — identical code path):

* **dev (default):** local file stores under ``CF_DATA_DIR`` (sqlite relational,
  lancedb vector, cognee's default graph). Chat via the Hugging Face router
  (``custom`` provider). Embeddings via local ``fastembed`` (no embeddings API).
* **prod:** set ``GRAPH_DATABASE_PROVIDER=neo4j`` + ``DB_PROVIDER=postgres`` +
  ``VECTOR_DB_PROVIDER=pgvector`` and the matching connection vars to externalize
  durable storage to Supabase + Neo4j Aura. The HF Space has no DB-wire firewall,
  so these connections work there (they cannot be exercised from the build
  sandbox — see project notes).

The build sandbox firewall blocks Postgres/Bolt wire protocols, so the prod
store branch is written here but is verified on the actual Space deploy.
"""
from __future__ import annotations

import os
from functools import lru_cache
from pathlib import Path
from typing import Any

from dotenv import load_dotenv

# .../backend/app/cognee_runtime/bootstrap.py -> parents[2] == .../backend
BACKEND_DIR = Path(__file__).resolve().parents[2]
ENV_PATH = BACKEND_DIR / ".env"


def _install_hf_router_embeddings(*, model, dimensions, api_key, endpoint, batch_size):
    """Route cognee's embedding resolution to our HF-router engine.

    cognee has no native adapter for the HF feature-extraction endpoint, so we
    replace the factory ``create_embedding_engine`` in its DEFINING module.

    Critically, we must reach the real *module* via ``importlib``: the
    ``embeddings`` package ``__init__`` does ``from .get_embedding_engine import
    get_embedding_engine``, which rebinds the package attribute
    ``...embeddings.get_embedding_engine`` to the *function* and shadows the
    submodule of the same name. So ``import ...get_embedding_engine as m`` binds
    ``m`` to the function, and setting ``m.create_embedding_engine`` is a silent
    no-op on a function object (this was the original bug: every call fell through
    to LiteLLM). ``get_embedding_engine()`` resolves ``create_embedding_engine``
    from its module globals at call time, so patching the module global routes
    every caller — chunking, vector-engine creation, search — to our engine.
    """
    import importlib

    from .hf_embeddings import HFRouterEmbeddingEngine

    engine = HFRouterEmbeddingEngine(
        model=model,
        dimensions=dimensions,
        api_key=api_key,
        endpoint=endpoint,
        batch_size=batch_size,
    )

    gee_mod = importlib.import_module(
        "cognee.infrastructure.databases.vector.embeddings.get_embedding_engine"
    )
    # Drop the original lru_cache (if any), then replace the factory wholesale.
    try:
        gee_mod.create_embedding_engine.cache_clear()
    except Exception:  # noqa: BLE001
        pass
    gee_mod.create_embedding_engine = lambda *a, **k: engine

    # Defensive: if a vector engine was already built+cached with the default
    # engine, evict it so the adapter is rebuilt holding our engine.
    try:
        cve = importlib.import_module(
            "cognee.infrastructure.databases.vector.create_vector_engine"
        )
        inner = getattr(cve, "_create_vector_engine", None)
        for attr in ("cache_clear", "cache_purge", "clear_cache"):
            fn = getattr(inner, attr, None)
            if callable(fn):
                fn()
                break
    except Exception:  # noqa: BLE001
        pass

    return engine


@lru_cache(maxsize=1)
def configure_cognee() -> dict[str, Any]:
    """Idempotently configure the Cognee SDK for the active profile.

    Returns a secret-free dict describing the resolved configuration (safe to log
    or return from a health endpoint).
    """
    load_dotenv(ENV_PATH, override=False)

    # --- keep cognee's stores inside the project (dev) / a writable mount (prod) ---
    data_dir = Path(os.getenv("CF_DATA_DIR", str(BACKEND_DIR / ".cf_data"))).resolve()
    system_root = data_dir / "system"
    data_root = data_dir / "data"
    system_root.mkdir(parents=True, exist_ok=True)
    data_root.mkdir(parents=True, exist_ok=True)

    # --- posture: single-tenant tool, no auth gate, no session cache by default ---
    os.environ.setdefault("ENABLE_BACKEND_ACCESS_CONTROL", "false")
    os.environ.setdefault("CACHING", "false")
    # Skip cognee's 30s embedding pre-flight: our HF-router engine has its own retry
    # path for cold model loads, and the pre-flight false-fails on a cold start.
    os.environ.setdefault("COGNEE_SKIP_CONNECTION_TEST", "true")

    import cognee

    cognee.config.system_root_directory(str(system_root))
    cognee.config.data_root_directory(str(data_root))

    # --- LLM: chat / graph extraction (cognify) ---
    llm_provider = os.getenv("LLM_PROVIDER", "custom")
    # Pin the HF router inference provider (":novita") — the bare model lets the
    # router auto-pick a provider, which intermittently returns 403 (e.g. deepinfra
    # not available for this key). novita serves Qwen2.5-72B + Llama-3.1-8B reliably.
    llm_model = os.getenv("LLM_MODEL", "openai/Qwen/Qwen2.5-72B-Instruct:novita")
    llm_endpoint = os.getenv("LLM_ENDPOINT", "https://router.huggingface.co/v1")
    llm_api_key = os.getenv("LLM_API_KEY") or os.getenv("HUGGINGFACE_API_KEY") or ""
    cognee.config.set_llm_config(
        {
            "llm_provider": llm_provider,
            "llm_model": llm_model,
            "llm_endpoint": llm_endpoint,
            "llm_api_key": llm_api_key,
        }
    )
    # HF router providers reject response_format=json_object (instructor's default
    # "json_mode"); markdown-JSON works with any OpenAI-compatible chat endpoint.
    # Default the custom/HF provider to markdown_json_mode.
    instructor_mode = os.getenv("LLM_INSTRUCTOR_MODE") or (
        "markdown_json_mode" if llm_provider == "custom" else ""
    )
    if instructor_mode:
        cognee.config.set_llm_config({"llm_instructor_mode": instructor_mode})

    # --- Embeddings ---
    # Default: "hf_router" -> call HF's feature-extraction API (no local model in RAM;
    # light enough for the 4 GB sandbox and a stateless Space). cognee has no native
    # adapter for that endpoint, so we inject a custom EmbeddingEngine by replacing the
    # factory `create_embedding_engine` (which get_embedding_engine() resolves from its
    # module globals at call time, so every caller is routed to our engine). Other
    # providers (fastembed, openai, openai_compatible, ollama) fall through to cognee.
    embedding_provider = os.getenv("EMBEDDING_PROVIDER", "hf_router")
    embedding_model = os.getenv("EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5")
    embedding_dimensions = int(os.getenv("EMBEDDING_DIMENSIONS", "384"))
    cognee.config.set_embedding_config(
        {
            "embedding_provider": embedding_provider,
            "embedding_model": embedding_model,
            "embedding_dimensions": embedding_dimensions,
        }
    )
    if embedding_provider == "hf_router":
        _install_hf_router_embeddings(
            model=embedding_model,
            dimensions=embedding_dimensions,
            api_key=os.getenv("EMBEDDING_API_KEY") or llm_api_key,
            endpoint=os.getenv("EMBEDDING_ENDPOINT") or None,
            batch_size=int(os.getenv("EMBEDDING_BATCH_SIZE", "16")),
        )

    # --- Graph store (prod: neo4j; dev: cognee default unless overridden) ---
    graph_provider = os.getenv("GRAPH_DATABASE_PROVIDER")
    if graph_provider:
        cognee.config.set_graph_database_provider(graph_provider)
        if graph_provider == "neo4j":
            cognee.config.set_graph_db_config(
                {
                    "graph_database_provider": "neo4j",
                    "graph_database_url": os.getenv("GRAPH_DATABASE_URL", ""),
                    "graph_database_username": os.getenv("GRAPH_DATABASE_USERNAME", "neo4j"),
                    "graph_database_password": os.getenv("GRAPH_DATABASE_PASSWORD", ""),
                }
            )

    # --- Relational store (prod: postgres / Supabase; dev: sqlite default) ---
    db_provider = os.getenv("DB_PROVIDER")
    if db_provider == "postgres":
        cognee.config.set_relational_db_config(
            {
                "db_provider": "postgres",
                "db_host": os.getenv("DB_HOST", ""),
                "db_port": os.getenv("DB_PORT", "5432"),
                "db_name": os.getenv("DB_NAME", "postgres"),
                "db_username": os.getenv("DB_USERNAME", ""),
                "db_password": os.getenv("DB_PASSWORD", ""),
            }
        )

    # --- Vector store (prod: pgvector on Supabase; dev: lancedb default) ---
    vector_provider = os.getenv("VECTOR_DB_PROVIDER")
    if vector_provider:
        cognee.config.set_vector_db_provider(vector_provider)

    return {
        "profile": "prod" if (db_provider or graph_provider == "neo4j") else "dev",
        "llm_provider": llm_provider,
        "llm_model": llm_model,
        "llm_endpoint": llm_endpoint,
        "llm_api_key_set": bool(llm_api_key),
        "embedding_provider": embedding_provider,
        "embedding_model": embedding_model,
        "embedding_dimensions": embedding_dimensions,
        "graph_provider": graph_provider or "cognee-default",
        "relational_provider": db_provider or "sqlite-default",
        "vector_provider": vector_provider or "lancedb-default",
        "data_dir": str(data_dir),
    }