Cognitive-rag / backend / config.py
config.py
Raw
from __future__ import annotations

import importlib
import importlib.util
import os
import sys
from pathlib import Path
from threading import Lock
from types import ModuleType
from typing import Any

from dotenv import load_dotenv


PROJECT_ROOT = Path(__file__).resolve().parent.parent
BACKEND_DIR = PROJECT_ROOT / "backend"
BACKEND_ENV_PATH = BACKEND_DIR / ".env"
BACKEND_ENV_EXAMPLE_PATH = BACKEND_DIR / ".env.example"

JOBS_DIR = BACKEND_DIR / "jobs"
AUDIT_LOG_PATH = BACKEND_DIR / "audit.log"
USER_REGISTRY_PATH = BACKEND_DIR / "user_registry.json"
RUN_REPORT_DIR = BACKEND_DIR / "run_report"
RUN_REPORT_UPLOAD_DIR = RUN_REPORT_DIR / "upload"
RUN_REPORT_PROMPT_DIR = RUN_REPORT_DIR / "prompt"

_BOOTSTRAPPED = False
_BOOTSTRAP_LOCK = Lock()
_VDB_MODULES: dict[str, ModuleType] = {}


class StartupConfigError(RuntimeError):
    pass


def _require_env_file() -> None:
    if BACKEND_ENV_PATH.exists():
        return
    raise StartupConfigError(
        "Missing backend/.env. Copy backend/.env.example to backend/.env and set required keys."
    )


def _load_backend_env() -> None:
    _require_env_file()
    load_dotenv(BACKEND_ENV_PATH, override=False)


def _ensure_runtime_paths() -> None:
    BACKEND_DIR.mkdir(parents=True, exist_ok=True)
    JOBS_DIR.mkdir(parents=True, exist_ok=True)
    RUN_REPORT_UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
    RUN_REPORT_PROMPT_DIR.mkdir(parents=True, exist_ok=True)
    if not AUDIT_LOG_PATH.exists():
        AUDIT_LOG_PATH.touch()
    if not USER_REGISTRY_PATH.exists():
        USER_REGISTRY_PATH.write_text("{}", encoding="utf-8")


def _put_stage_paths_on_sys_path() -> None:
    stage_paths = [
        PROJECT_ROOT / "01_data_ingestion",
        PROJECT_ROOT / "02_chunking",
        PROJECT_ROOT / "03_embedding",
        PROJECT_ROOT / "04_vectoredb",
        PROJECT_ROOT / "05_generation",
        PROJECT_ROOT / "06_conversation_memory",
    ]
    for path in stage_paths:
        value = str(path)
        if value not in sys.path:
            sys.path.insert(0, value)


def _require_env_value(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        raise StartupConfigError(
            f"Missing required env var: {name}. Update backend/.env."
        )
    return value


def _load_vectordb_config_module() -> ModuleType:
    vectordb_config_path = PROJECT_ROOT / "04_vectoredb" / "config.py"
    try:
        spec = importlib.util.spec_from_file_location("config", vectordb_config_path)
        if spec is None or spec.loader is None:
            raise RuntimeError("Unable to load 04_vectoredb/config.py module spec.")
        module = importlib.util.module_from_spec(spec)
        # Ensure retriever/pinecone modules import this exact config module.
        sys.modules["config"] = module
        spec.loader.exec_module(module)
    except Exception as exc:
        message = str(exc).lower()
        if ".envpinecone" in message or "envpinecone" in message:
            raise StartupConfigError(
                "Failed to import 04_vectoredb/config.py due to .envpinecone dependency. "
                "This backend will not create files outside backend/. "
                "Next steps: either make upstream vectordb config env-compatible, or manually "
                "provision the expected 04_vectoredb/.envpinecone outside this task scope."
            ) from exc
        raise StartupConfigError(
            f"Failed to import 04_vectoredb/config.py: {exc}"
        ) from exc

    # Patch in-memory values before importing pinecone/retriever modules.
    module.PINECONE_API_KEY = _require_env_value("PINECONE_API_KEY")
    for env_name, attr_name in (
        ("PINECONE_INDEX_NAME", "PINECONE_INDEX_NAME"),
        ("SPARSE_INDEX_NAME", "SPARSE_INDEX_NAME"),
        ("PINECONE_CLOUD", "PINECONE_CLOUD"),
        ("PINECONE_REGION", "PINECONE_REGION"),
        ("TOP_K", "TOP_K"),
        ("RRF_K", "RRF_K"),
        ("RERANKER_ENABLED", "RERANKER_ENABLED"),
        ("RERANKER_MODEL", "RERANKER_MODEL"),
        ("RERANKER_TOP_N", "RERANKER_TOP_N"),
        ("MEMORY_BOOST", "MEMORY_BOOST"),
        ("MEMORY_MIN_SCORE", "MEMORY_MIN_SCORE"),
        ("LANGCHAIN_API_KEY", "LANGCHAIN_API_KEY"),
        ("LANGCHAIN_PROJECT", "LANGCHAIN_PROJECT"),
        ("LANGCHAIN_TRACING_V2", "LANGCHAIN_TRACING_V2"),
        ("NVIDIA_API_KEY", "NVIDIA_API_KEY"),
        ("EMBEDDING_MODEL", "EMBEDDING_MODEL"),
    ):
        env_value = os.environ.get(env_name)
        if env_value is None:
            continue
        if attr_name in {"TOP_K", "RRF_K", "RERANKER_TOP_N"}:
            setattr(module, attr_name, int(env_value))
        elif attr_name in {"MEMORY_BOOST", "MEMORY_MIN_SCORE"}:
            setattr(module, attr_name, float(env_value))
        elif attr_name == "RERANKER_ENABLED":
            setattr(module, attr_name, env_value.lower() in ("true", "1", "yes"))
        else:
            setattr(module, attr_name, env_value)

    return module


def bootstrap() -> None:
    global _BOOTSTRAPPED
    with _BOOTSTRAP_LOCK:
        if _BOOTSTRAPPED:
            return
        _load_backend_env()
        _ensure_runtime_paths()
        _put_stage_paths_on_sys_path()
        _load_vectordb_config_module()
        _BOOTSTRAPPED = True


def _import_module(name: str) -> ModuleType:
    module = importlib.import_module(name)
    return module


def get_stage_modules() -> dict[str, ModuleType]:
    """
    Lazy imports after bootstrap.
    """
    bootstrap()
    if _VDB_MODULES:
        return _VDB_MODULES

    modules: dict[str, ModuleType] = {}
    modules["ingestion_engine"] = _import_module("ingestion.engine")
    modules["ingestion_events"] = _import_module("ingestion.events")
    modules["chunking"] = _import_module("chunk_documents")
    modules["embedding_chunks"] = _import_module("embed_chunks")
    modules["embedding_query"] = _import_module("embed_query")
    modules["vectordb_config"] = _import_module("config")
    modules["vectordb_pinecone"] = _import_module("pinecone_client")
    modules["vectordb_upsert"] = _import_module("upsert_db")
    modules["vectordb_retriever"] = _import_module("retriever_multi")
    modules["generation"] = _import_module("generate_answer")
    modules["memory"] = _import_module("memory_processor")
    _VDB_MODULES.update(modules)
    return modules


def env_value(name: str, default: str = "") -> str:
    return os.environ.get(name, default)


def env_int(name: str, default: int) -> int:
    value = os.environ.get(name)
    if value is None:
        return default
    return int(value)


def job_dir(job_id: str) -> Path:
    return JOBS_DIR / job_id


def job_uploads_dir(job_id: str) -> Path:
    return job_dir(job_id) / "uploads"


def job_ingestion_dir(job_id: str) -> Path:
    return job_dir(job_id) / "ingestion"


def job_chunking_dir(job_id: str) -> Path:
    return job_dir(job_id) / "chunking"


def job_embedding_dir(job_id: str) -> Path:
    return job_dir(job_id) / "embedding"


def response_source_from_chunk(chunk: dict[str, Any]) -> dict[str, Any]:
    score = chunk.get("cross_encoder_score")
    if score is None:
        score = chunk.get("score")
    if score is None:
        score = chunk.get("rrf_score")

    output: dict[str, Any] = {}
    for key in ("chunk_id", "doc_id"):
        if chunk.get(key) is not None and chunk.get(key) != "":
            output[key] = chunk.get(key)
    if score is not None:
        output["score"] = score
    for key in ("page", "section"):
        value = chunk.get(key)
        if value is not None and value != "":
            output[key] = value
    source_path = chunk.get("source_path")
    if source_path:
        output["source_path"] = source_path
    return output