Agentic-OS/backend/app/services/model_config.py

"""Runtime LLM routing configuration (env defaults + Redis overrides)."""
from __future__ import annotations

import json
import time
from dataclasses import asdict, dataclass, field

from app.config import settings
from app.services.model_catalog import ModelEntry, entry_for_model, get_entry

ROUTING_KEY = "agentic:llm:routing"

# Short in-process cache: routing config is read many times per task but
# changes only via the admin /api/llm/routing endpoint.
_CONFIG_CACHE_TTL = 30.0
_config_cache: tuple[float, "LlmRoutingConfig"] | None = None

DEFAULT_CLOUD_PROVIDER_ORDER = ["gemini", "deepseek", "openrouter"]


@dataclass
class LlmRoutingConfig:
    """Selected models per tier. Always tries local first, then escalates."""

    local_model_id: str = "local-qwen25-7b"
    economy_gemini_id: str = "gemini-flash-25"
    economy_deepseek_id: str = "ds-chat"
    economy_openrouter_id: str = "or-ds-chat"
    premium_gemini_id: str = "gemini-pro-25"
    premium_deepseek_id: str = "ds-reasoner"
    premium_openrouter_id: str = "or-claude-35-sonnet"
    # Cloud provider cascade within each tier: gemini → deepseek → openrouter
    cloud_provider_order: list[str] = field(default_factory=lambda: list(DEFAULT_CLOUD_PROVIDER_ORDER))
    auto_escalate: bool = True
    # Cap escalation: local | economy | premium
    max_tier: str = "premium"

    def to_dict(self) -> dict:
        return asdict(self)

    def resolve(self, entry_id: str) -> ModelEntry | None:
        entry = get_entry(entry_id)
        if entry:
            return entry
        return None


def _defaults_from_env() -> LlmRoutingConfig:
    """Build defaults using env vars and catalog IDs where possible."""
    local = settings.llm_local_model
    cfg = LlmRoutingConfig(
        local_model_id=_find_id("ollama", local, "local") or "local-qwen25-7b",
        economy_gemini_id=_find_id("gemini", settings.llm_economy_gemini, "economy") or "gemini-flash-25",
        economy_deepseek_id=_find_id("deepseek", settings.llm_economy_deepseek, "economy") or "ds-chat",
        economy_openrouter_id=_find_id("openrouter", settings.llm_economy_openrouter, "economy") or "or-ds-chat",
        premium_gemini_id=_find_id("gemini", settings.llm_premium_gemini, "premium") or "gemini-pro-25",
        premium_deepseek_id=_find_id("deepseek", settings.llm_premium_deepseek, "premium") or "ds-reasoner",
        premium_openrouter_id=_find_id("openrouter", settings.llm_premium_openrouter, "premium")
        or "or-claude-35-sonnet",
        auto_escalate=settings.llm_auto_escalate,
        max_tier=settings.llm_max_tier,
    )
    if settings.llm_cloud_provider_order:
        cfg.cloud_provider_order = [
            p.strip() for p in settings.llm_cloud_provider_order.split(",") if p.strip()
        ]
    return cfg


def _find_id(backend: str, model: str, tier: str) -> str | None:
    from app.services.model_catalog import CATALOG

    for e in CATALOG:
        if e.backend == backend and e.model == model and e.tier == tier:
            return e.id
    return None


async def get_routing_config(*, use_cache: bool = True) -> LlmRoutingConfig:
    from app.services.events import get_redis

    global _config_cache
    if use_cache and _config_cache is not None:
        ts, cached = _config_cache
        if time.monotonic() - ts < _CONFIG_CACHE_TTL:
            return cached

    cfg = _defaults_from_env()
    try:
        raw = await get_redis().get(ROUTING_KEY)
        if raw:
            data = json.loads(raw)
            for k, v in data.items():
                if hasattr(cfg, k):
                    setattr(cfg, k, v)
    except Exception:  # noqa: BLE001
        pass

    _config_cache = (time.monotonic(), cfg)
    return cfg


async def save_routing_config(cfg: LlmRoutingConfig) -> None:
    from app.services.events import get_redis

    global _config_cache
    await get_redis().set(ROUTING_KEY, json.dumps(cfg.to_dict()))
    _config_cache = (time.monotonic(), cfg)


def available_backends() -> dict[str, bool]:
    return {
        "ollama": True,
        "gemini": bool(settings.gemini_api_key),
        "deepseek": bool(settings.deepseek_api_key),
        "openrouter": bool(settings.openrouter_api_key),
    }


def _provider_entry(cfg: LlmRoutingConfig, tier: str, provider: str) -> ModelEntry | None:
    avail = available_backends()
    if tier == "economy":
        if provider == "gemini" and avail["gemini"]:
            return cfg.resolve(cfg.economy_gemini_id) or entry_for_model(
                "gemini", settings.llm_economy_gemini, "economy"
            )
        if provider == "deepseek" and avail["deepseek"]:
            return cfg.resolve(cfg.economy_deepseek_id) or entry_for_model(
                "deepseek", settings.llm_economy_deepseek, "economy"
            )
        if provider == "openrouter" and avail["openrouter"]:
            return cfg.resolve(cfg.economy_openrouter_id) or entry_for_model(
                "openrouter", settings.llm_economy_openrouter, "economy"
            )
    elif tier == "premium":
        if provider == "gemini" and avail["gemini"]:
            return cfg.resolve(cfg.premium_gemini_id) or entry_for_model(
                "gemini", settings.llm_premium_gemini, "premium"
            )
        if provider == "deepseek" and avail["deepseek"]:
            return cfg.resolve(cfg.premium_deepseek_id) or entry_for_model(
                "deepseek", settings.llm_premium_deepseek, "premium"
            )
        if provider == "openrouter" and avail["openrouter"]:
            return cfg.resolve(cfg.premium_openrouter_id) or entry_for_model(
                "openrouter", settings.llm_premium_openrouter, "premium"
            )
    return None


def resolve_tier_models(cfg: LlmRoutingConfig, tier: str) -> list[ModelEntry]:
    """Return ordered models for a tier (local = one entry; cloud = provider cascade)."""
    if tier == "local":
        if not available_backends()["ollama"]:
            return []
        e = cfg.resolve(cfg.local_model_id) or entry_for_model("ollama", settings.llm_local_model, "local")
        return [e]

    models: list[ModelEntry] = []
    seen: set[str] = set()
    for provider in cfg.cloud_provider_order:
        entry = _provider_entry(cfg, tier, provider)
        if entry and entry.id not in seen:
            models.append(entry)
            seen.add(entry.id)

    # Fallback: any configured cloud provider with a key
    for provider in DEFAULT_CLOUD_PROVIDER_ORDER:
        if provider in cfg.cloud_provider_order:
            continue
        entry = _provider_entry(cfg, tier, provider)
        if entry and entry.id not in seen:
            models.append(entry)
            seen.add(entry.id)
    return models


def resolve_tier_model(cfg: LlmRoutingConfig, tier: str) -> ModelEntry | None:
    """Pick the first available model for a tier (legacy single-model callers)."""
    models = resolve_tier_models(cfg, tier)
    return models[0] if models else None