pablan/backend/app/llm/overrides.py

"""The LLM endpoint configuration, as the process sees it.

**Bootstrap, then DB.** On the very first start the `PABLAN_*` environment
variables are copied into `llm_settings`, one row per role. From that
moment the table is the truth: later `.env` edits are ignored, because a
configuration an admin can change in the UI and a configuration the
deployment can change underneath them cannot both be authoritative. The
environment stays reachable as the value a field can be *reset* to, which
is what `env_defaults()` is for.

The configuration lives in a module-level cache because `_role_config` is a
hot, synchronous function on every LLM call — it cannot await a query. The
cache is filled at startup and refreshed whenever an admin writes, which is
also when the OpenAI clients are rebuilt.

Single-process by design: the customer stack pins `--workers 1` (see
architecture.md), so there is exactly one cache to refresh. A multi-worker
deployment would need a notification channel instead.
"""

import logging
from dataclasses import dataclass

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import get_settings
from app.models import LLMSetting

logger = logging.getLogger("pablan.llm")

ROLES = ("chat", "utility", "embedding")


@dataclass(frozen=True)
class RoleConfig:
    """One role's stored configuration. A None field means the column is
    empty, which after bootstrap only happens if an admin blanked it."""

    base_url: str | None = None
    model: str | None = None
    api_key: str | None = None


_config: dict[str, RoleConfig] = {}


def env_defaults(role: str) -> RoleConfig:
    """What `.env` says for this role — the value "reset to .env" restores.

    Read live rather than remembered from bootstrap: an admin who fixes a
    typo in `.env` and resets the field should get the corrected value, not
    the one that was wrong at install time.
    """
    settings = get_settings()
    base_url, api_key, model = {
        "chat": (settings.chat_base_url, settings.chat_api_key, settings.chat_model),
        "utility": (
            settings.utility_base_url,
            settings.utility_api_key,
            settings.utility_model,
        ),
        "embedding": (
            settings.embedding_base_url,
            settings.embedding_api_key,
            settings.embedding_model,
        ),
    }[role]
    return RoleConfig(base_url=base_url, model=model, api_key=api_key)


_FIELDS = ("base_url", "model", "api_key")


async def bootstrap_llm_settings(db: AsyncSession) -> int:
    """Copy the environment into any field that still defers to it.

    Runs at every startup, but only ever fills blanks: a field is written
    exactly when it is flagged `*_from_env` AND currently empty. That is
    true for a fresh install (no rows yet) and for a field an upgrade
    marked as still belonging to `.env`, and false for anything an admin
    has typed, which is never touched.

    Returns the number of fields written.
    """
    rows = {row.role: row for row in (await db.execute(select(LLMSetting))).scalars()}
    written = 0
    for role in ROLES:
        row = rows.get(role)
        if row is None:
            # Flags set explicitly rather than left to the column defaults:
            # those only materialise on flush, and the loop below reads them
            # before that.
            row = LLMSetting(
                role=role,
                base_url_from_env=True,
                model_from_env=True,
                api_key_from_env=True,
            )
            db.add(row)
        defaults = env_defaults(role)
        for field in _FIELDS:
            if not getattr(row, f"{field}_from_env") or getattr(row, field):
                continue
            setattr(row, field, getattr(defaults, field) or None)
            written += 1
    if written:
        await db.commit()
    # Counts and roles only — never the values, one of which is a key.
    # NB: not `created` — logging reserves that name on
    # LogRecord and raises KeyError when an `extra` key collides with it.
    logger.info(
        "llm settings bootstrapped",
        extra={"event": "llm_bootstrap", "fields_written": written},
    )
    return written


async def load_config(db: AsyncSession) -> None:
    """Re-read every stored row. Call after any write."""
    rows = (await db.execute(select(LLMSetting))).scalars().all()
    _config.clear()
    _config.update(
        {
            row.role: RoleConfig(
                base_url=row.base_url, model=row.model, api_key=row.api_key
            )
            for row in rows
        }
    )
    logger.info(
        "llm settings loaded",
        extra={"event": "llm_settings_loaded", "roles": sorted(_config)},
    )


def get_config(role: str) -> RoleConfig:
    return _config.get(role, RoleConfig())


def clear() -> None:
    """Drop the cache — used by tests between cases."""
    _config.clear()