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

"""Tiered model routing: local first, auto-escalate to economy/premium when needed."""
from __future__ import annotations

import logging
from typing import Any

from langchain_core.messages import BaseMessage
from langchain_openai import ChatOpenAI

from app.config import settings
from app.services.llm_usage import LlmUsageTracker, ModelInfo
from app.services.model_catalog import ModelEntry
from app.services.model_config import LlmRoutingConfig, get_routing_config, resolve_tier_models

logger = logging.getLogger(__name__)

TIER_ORDER = ["local", "economy", "premium"]
TIER_RANK = {t: i for i, t in enumerate(TIER_ORDER)}

GEMINI_OPENAI_BASE = "https://generativelanguage.googleapis.com/v1beta/openai/"


def build_chat_model(entry: ModelEntry, temperature: float = 0.1) -> ChatOpenAI:
    if entry.backend == "ollama":
        return ChatOpenAI(
            model=entry.model,
            temperature=temperature,
            base_url=f"{settings.ollama_base_url.rstrip('/')}/v1",
            api_key="ollama",
            timeout=120,
            max_retries=2,
        )
    if entry.backend == "gemini":
        return ChatOpenAI(
            model=entry.model,
            temperature=temperature,
            base_url=GEMINI_OPENAI_BASE,
            api_key=settings.gemini_api_key,
            timeout=120,
            max_retries=2,
        )
    if entry.backend == "deepseek":
        return ChatOpenAI(
            model=entry.model,
            temperature=temperature,
            base_url="https://api.deepseek.com/v1",
            api_key=settings.deepseek_api_key,
            timeout=120,
            max_retries=2,
        )
    return ChatOpenAI(
        model=entry.model,
        temperature=temperature,
        base_url="https://openrouter.ai/api/v1",
        api_key=settings.openrouter_api_key,
        timeout=120,
        max_retries=2,
        default_headers={"HTTP-Referer": "http://localhost:3000", "X-Title": "Agentic OS"},
    )


def entry_to_info(entry: ModelEntry, step: str) -> ModelInfo:
    return ModelInfo(
        step=step,
        provider=entry.provider,
        backend=entry.backend,
        model=entry.model,
        display=entry.display,
    )


def plan_reasoning_tiers(
    triage: dict,
    diagnostics: list[dict],
    *,
    prior_context: dict | None = None,
    follow_up: str | None = None,
    run_number: int = 1,
    cfg: LlmRoutingConfig | None = None,
    min_tier: str | None = None,
) -> tuple[list[str], str]:
    """Return ordered tiers to try and a human-readable routing rationale."""
    cfg = cfg or LlmRoutingConfig()
    severity = (triage.get("severity") or "medium").lower()
    recommended = (triage.get("recommended_tier") or "local").lower()
    needs_cloud = bool(triage.get("needs_cloud_reasoning"))
    fail_count = sum(1 for d in diagnostics if not d.get("ok"))
    diag_count = len(diagnostics)

    max_tier = "local"
    reasons: list[str] = ["Always start with local LLM"]

    if needs_cloud:
        max_tier = "economy"
        reasons.append("Triage flagged cloud reasoning helpful")
    if severity in ("high", "critical"):
        max_tier = "premium" if severity == "critical" else max([max_tier, "economy"], key=lambda t: TIER_RANK[t])
        reasons.append(f"Severity is {severity}")
    if recommended == "premium":
        max_tier = "premium"
        reasons.append("Triage recommended premium tier")
    elif recommended == "economy" and TIER_RANK[max_tier] < TIER_RANK["economy"]:
        max_tier = "economy"
        reasons.append("Triage recommended economy tier")
    if fail_count >= 2 or (fail_count >= 1 and diag_count >= 4 and fail_count / max(diag_count, 1) > 0.25):
        max_tier = max([max_tier, "economy"], key=lambda t: TIER_RANK[t])
        reasons.append(f"{fail_count} diagnostic failure(s)")
    if fail_count >= 4 or (severity == "critical" and fail_count >= 1):
        max_tier = "premium"
        reasons.append("Complex failure pattern — premium model")
    if prior_context and not prior_context.get("resolved"):
        max_tier = max([max_tier, "economy"], key=lambda t: TIER_RANK[t])
        reasons.append("Prior run unresolved")
    if follow_up:
        max_tier = max([max_tier, "economy"], key=lambda t: TIER_RANK[t])
        reasons.append("Follow-up investigation")
    if run_number > 1:
        max_tier = max([max_tier, "economy"], key=lambda t: TIER_RANK[t])
        reasons.append(f"Run {run_number} continuation")

    # Respect configured cap
    cap = cfg.max_tier if cfg.max_tier in TIER_RANK else "premium"
    if TIER_RANK[max_tier] > TIER_RANK[cap]:
        max_tier = cap
        reasons.append(f"Capped at {cap} tier")

    max_idx = TIER_RANK[max_tier]
    tiers = TIER_ORDER[: max_idx + 1]

    if not cfg.auto_escalate and len(tiers) > 1:
        tiers = ["local"]
        reasons.append("Auto-escalation disabled — local only")

    if min_tier and min_tier in TIER_RANK:
        from app.services.task_runtime import apply_min_tier_floor

        tiers = apply_min_tier_floor(tiers, min_tier)
        reasons.append(f"User requested {min_tier} API (skip local)")

    return tiers, "; ".join(reasons)


def _should_escalate(
    analysis: dict,
    tier: str,
    tiers: list[str],
    *,
    diag_ok_ratio: float = 0.0,
    is_last_provider: bool = False,
) -> bool:
    if is_last_provider and tier == tiers[-1]:
        return False
    if analysis.get("error"):
        return True
    confidence = analysis.get("confidence")
    if analysis.get("needs_escalation"):
        return True
    if confidence is not None and float(confidence) < 0.65 and diag_ok_ratio < 0.75:
        return True
    if not analysis.get("summary") and not analysis.get("root_cause"):
        return True
    return False


async def reason_with_cascade(
    tracker: LlmUsageTracker,
    tiers: list[str],
    messages: list[BaseMessage],
    parse_json,
    cfg: LlmRoutingConfig | None = None,
    *,
    diagnostics: list[dict] | None = None,
    task_id: int | None = None,
) -> tuple[dict, dict]:
    """Try tiers in order; within cloud tiers cascade gemini → deepseek → openrouter."""
    cfg = cfg or await get_routing_config()
    analysis: dict[str, Any] = {}
    used_tiers: list[str] = []
    used_backends: list[str] = []
    last_entry: ModelEntry | None = None
    last_error: str | None = None
    diag_ok_ratio = 0.0
    if diagnostics:
        diag_ok_ratio = sum(1 for d in diagnostics if d.get("ok")) / len(diagnostics)

    from app.services.task_runtime import get_task_min_tier

    finished = False
    for tier in tiers:
        min_tier = await get_task_min_tier(task_id) if task_id else None
        if min_tier and TIER_RANK[tier] < TIER_RANK[min_tier]:
            continue

        entries = resolve_tier_models(cfg, tier)
        if not entries:
            logger.info("No models available for tier %s", tier)
            continue

        for idx, entry in enumerate(entries):
            min_tier = await get_task_min_tier(task_id) if task_id else None
            if min_tier and TIER_RANK[tier] < TIER_RANK[min_tier]:
                break

            last_entry = entry
            is_last_provider = idx == len(entries) - 1
            step_label = f"reason-{tier}-{entry.backend}"
            try:
                model = build_chat_model(entry, temperature=0.1)
                info = entry_to_info(entry, step_label)
                resp = await tracker.invoke(step_label, model, info, messages)
                analysis = parse_json(resp.content) or {
                    "summary": str(resp.content)[:1000],
                    "resolved": False,
                }
                analysis["_tier_used"] = tier
                analysis["_backend_used"] = entry.backend
                analysis["_model"] = entry.display
                used_tiers.append(tier)
                used_backends.append(entry.backend)

                if not cfg.auto_escalate or not _should_escalate(
                    analysis,
                    tier,
                    tiers,
                    diag_ok_ratio=diag_ok_ratio,
                    is_last_provider=is_last_provider,
                ):
                    min_tier = await get_task_min_tier(task_id) if task_id else None
                    if not (
                        min_tier
                        and TIER_RANK.get(min_tier, 0) > TIER_RANK.get(tier, 0)
                    ):
                        finished = True
                        break
                logger.info(
                    "Escalating from %s/%s: confidence=%s needs_esc=%s",
                    tier,
                    entry.backend,
                    analysis.get("confidence"),
                    analysis.get("needs_escalation"),
                )
            except Exception as exc:  # noqa: BLE001
                last_error = str(exc)
                logger.warning("Reasoning %s/%s failed: %s", tier, entry.backend, exc)
                analysis = {
                    "summary": f"reasoning error ({tier}/{entry.backend}): {exc}",
                    "resolved": False,
                    "error": str(exc),
                }
                used_tiers.append(tier)
                used_backends.append(entry.backend)
                continue

        if finished:
            break

    if last_error and not analysis.get("summary"):
        analysis = {"summary": f"reasoning error: {last_error}", "resolved": False, "error": last_error}

    routing_meta = {
        "tiers_planned": tiers,
        "tiers_used": used_tiers,
        "backends_used": used_backends,
        "final_tier": used_tiers[-1] if used_tiers else None,
        "final_backend": last_entry.backend if last_entry else None,
        "final_model": last_entry.display if last_entry else None,
    }
    return analysis, routing_meta