6krrt/tui_model.py

"""Pure data layer for the LLM Router TUI — no Textual dependency.

Keeps the payload-shaping logic (``build_model``, ``build_category_breakdown``)
and the HTTP fetcher out of ``tui.py`` so that module stays a thin rendering
shell, and so this layer is importable and testable without a running TUI.
"""

from __future__ import annotations

from collections import Counter, defaultdict
from typing import Any

import requests

DEFAULT_BASE_URL = "http://127.0.0.1:8080"

__all__ = [
    "fetch_metrics",
    "build_model",
    "build_category_breakdown",
    "decision_row",
]


def fetch_metrics(base_url: str) -> dict:
    """GET ``<base_url>/metrics`` and return the parsed JSON dict.

    Raises on any HTTP or network error (via ``raise_for_status`` and the
    underlying ``requests`` exception) so the caller can catch it and render
    a "cannot reach router" state.
    """
    url = base_url.rstrip("/") + "/metrics"
    resp = requests.get(url, timeout=10)
    resp.raise_for_status()
    return resp.json()


def build_model(data: dict) -> dict:
    """Turn raw /metrics JSON into a plain dict of rendered panel payloads.

    Returns keys: ``quota`` (list of {label, value} rows), ``per_model``
    (list of rows), ``verdict_mix`` (list of {verdict, count}),
    ``recent_decisions`` (list of rows), ``category_breakdown`` (list of
    rows), ``warnings`` (list of strings).
    """
    quota = data.get("quota")
    if quota:
        quota_rows = [
            {"label": "plan_kwh", "value": quota.get("plan_kwh")},
            {"label": "metered_kwh_30d", "value": quota.get("metered_kwh_30d")},
            {
                "label": "fraction",
                "value": quota.get("metered_fraction_of_plan"),
            },
            {"label": "calls", "value": quota.get("metered_calls_30d")},
        ]
    else:
        quota_rows = [{"label": "quota", "value": "N/A (plan not set)"}]

    per_model = [
        {
            "model": r.get("model_id"),
            "calls": r.get("calls"),
            "cost_usd": r.get("sum_cost_usd"),
            "energy_kwh": r.get("sum_energy_kwh"),
            "carbon_g_co2eq": r.get("sum_carbon_g_co2eq"),
        }
        for r in (data.get("per_model") or [])
    ]

    mix = data.get("verdict_mix") or {}
    verdict_mix = [
        {"verdict": k, "count": v} for k, v in sorted(mix.items())
    ]

    recent = []
    for r in data.get("recent_decisions") or []:
        recent.append(decision_row(r))

    coverage = data.get("coverage") or {}
    warnings = list(coverage.get("warnings") or [])

    return {
        "quota": quota_rows,
        "per_model": per_model,
        "verdict_mix": verdict_mix,
        "recent_decisions": recent,
        "category_breakdown": build_category_breakdown(recent),
        "warnings": warnings,
    }


def decision_row(r: dict) -> dict:
    """Project one route_decisions row onto the TUI's enriched decision shape.

    Shared by ``build_model`` (from /metrics) and the live SSE path so the two
    never drift in the fields they surface to the detail popup / table.
    """
    return {
        "id": r.get("id"),
        "kind": r.get("kind"),
        "category": r.get("task_category"),
        "tier": r.get("task_tier"),
        "selected": r.get("selected_model") or "none",
        "selected_provider": r.get("selected_provider"),
        "est_cost_usd": r.get("est_cost_usd"),
        "required_context_tokens": r.get("required_context_tokens"),
        "confidence": r.get("confidence"),
        "classifier_ms": r.get("classifier_ms"),
        "classification_source": r.get("classification_source"),
        "latency_tolerance": r.get("latency_tolerance"),
        "candidates_considered": r.get("candidates_considered"),
        "runner_up_models": r.get("runner_up_models"),
        "est_proficiency": r.get("est_proficiency"),
        "rejected_reason": r.get("rejected_reason"),
        "observed_at": r.get("observed_at"),
        "tools": r.get("tools"),
        "images": r.get("images"),
        "json_mode": r.get("json_mode"),
        "streamed": r.get("streamed"),
    }


def build_category_breakdown(
    decisions: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Aggregate recent decisions by (category, tier), with the most-common
    selected model and its share.

    Returns a list of dicts: ``category``, ``tier``, ``count``, ``majority``
    (the selected model with the most wins), ``share`` (its fraction of the
    count as a float in 0..1). Newer rows add first so ties settle toward the
    more recent model. This is the panel that answers "what is routing sending
    coding_general to right now?" without reading every individual decision.
    """
    ordered = sorted(
        decisions, key=lambda d: (d.get("id") or 0), reverse=True
    )
    buckets: dict[tuple[Any, Any], Counter] = defaultdict(Counter)
    counts: dict[tuple[Any, Any], int] = defaultdict(int)
    for decision in ordered:
        key = (decision.get("category"), decision.get("tier"))
        counts[key] += 1
        selected = decision.get("selected")
        if selected is not None:
            buckets[key][str(selected)] += 1

    rows = []
    for (category, tier), n in counts.items():
        winners = buckets[(category, tier)]
        if winners:
            majority, majority_count = winners.most_common(1)[0]
            share = majority_count / n
        else:
            majority, share = None, 0.0
        rows.append(
            {
                "category": category,
                "tier": tier,
                "count": n,
                "majority": majority if majority is not None else "none",
                "share": round(share, 2),
            }
        )
    return rows