app/route.py

"""Resolve Gemma 4 12B Unified: direct vLLM, Hermes custom_providers, or Studio API."""

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Callable
from urllib.parse import urlparse

import httpx

from app.config import Settings

GEMMA_MODEL_MARKERS = ("gemma-4-12b", "gemma4-12b", "gemma-4-12B", "12b-it", "12B-it")


@dataclass(frozen=True)
class Brain:
    kind: str  # direct | hermes | studio
    base_url: str
    model: str
    source: str


def _norm(url: str) -> str:
    url = url.strip().rstrip("/")
    if url and not url.endswith("/v1") and "/v1" not in urlparse(url).path:
        url = url + "/v1"
    return url


def looks_like_gemma(name: str) -> bool:
    lower = name.lower()
    return any(m.lower() in lower for m in GEMMA_MODEL_MARKERS) or "gemma" in lower


def parse_hermes_providers(text: str) -> list[Brain]:
    """Pull OpenAI-compat custom_providers from a Hermes config.yaml."""
    found: list[Brain] = []
    block = text.split("custom_providers:", 1)
    if len(block) < 2:
        return found
    body = block[1]
    chunks = re.split(r"\n  - name:", body)
    for chunk in chunks[1:]:
        name_m = re.match(r"\s*([^\n]+)", chunk)
        url_m = re.search(r"base_url:\s*(\S+)", chunk)
        if not url_m:
            continue
        name = (name_m.group(1).strip() if name_m else "")
        base = url_m.group(1).strip().strip("\"'")
        models = re.findall(r"\n      - (\S+)", chunk)
        if looks_like_gemma(name) or any(looks_like_gemma(m) for m in models):
            model = next((m for m in models if looks_like_gemma(m)), models[0] if models else "google/gemma-4-12B-it")
            found.append(
                Brain(kind="hermes", base_url=_norm(base), model=model, source=f"hermes:{name}")
            )
    return found


def brains_from_hermes_file(path: Path) -> list[Brain]:
    if not path.is_file():
        return []
    try:
        return parse_hermes_providers(path.read_text(encoding="utf-8"))
    except OSError:
        return []


def probe_models(base_url: str, *, timeout_s: float = 1.5, client: httpx.Client | None = None) -> list[str]:
    url = _norm(base_url)
    own = client is None
    http = client or httpx.Client(timeout=timeout_s)
    try:
        response = http.get(f"{url}/models")
        if response.status_code >= 500:
            return []
        payload = response.json()
        rows = payload.get("data", payload if isinstance(payload, list) else [])
        ids: list[str] = []
        for row in rows:
            if isinstance(row, dict) and row.get("id"):
                ids.append(str(row["id"]))
            elif isinstance(row, str):
                ids.append(row)
        return ids
    except (httpx.HTTPError, ValueError, TypeError):
        return []
    finally:
        if own:
            http.close()


def probe_studio(studio_url: str, *, timeout_s: float = 1.5, client: httpx.Client | None = None) -> bool:
    url = studio_url.rstrip("/")
    own = client is None
    http = client or httpx.Client(timeout=timeout_s)
    try:
        try:
            response = http.get(f"{url}/api/health")
            if response.status_code < 500:
                return True
        except httpx.HTTPError:
            pass
        try:
            response = http.get(f"{url}/phone")
            return response.status_code < 500
        except httpx.HTTPError:
            return False
    finally:
        if own:
            http.close()


ProbeFn = Callable[[str], list[str]]
StudioProbeFn = Callable[[str], bool]


def candidate_brains(settings: Settings) -> list[Brain]:
    out: list[Brain] = []
    seen: set[str] = set()

    def add(brain: Brain) -> None:
        key = f"{brain.kind}|{brain.base_url}|{brain.model}"
        if key not in seen:
            seen.add(key)
            out.append(brain)

    model = settings.llm_model
    add(Brain("direct", _norm(settings.llm_base_url), model, "RECEIPT_LLM_BASE_URL"))
    if settings.hermes_base_url:
        add(Brain("hermes", _norm(settings.hermes_base_url), settings.hermes_model or model, "RECEIPT_HERMES_BASE_URL"))
    for brain in brains_from_hermes_file(settings.hermes_config_path):
        add(brain)
    add(Brain("direct", "http://127.0.0.1:8080/v1", model, "localhost:8080"))
    host = (settings.gpu_host or "").strip()
    if host:
        add(Brain("direct", _norm(f"http://{host}:8080/v1"), model, f"RECEIPT_GPU_HOST:{host}"))
    return out


def resolve_brain(
    settings: Settings,
    *,
    probe: ProbeFn | None = None,
    studio_probe: StudioProbeFn | None = None,
) -> Brain:
    """Pick Gemma 4 12B: studio (Lamp) vs direct vLLM vs Hermes-discovered provider."""
    route = settings.llm_route.lower().strip()
    model = settings.llm_model

    if route == "studio":
        url = (settings.studio_url or "").rstrip("/")
        if not url:
            raise RuntimeError("RECEIPT_LLM_ROUTE=studio requires RECEIPT_STUDIO_URL")
        return Brain("studio", url, model, "RECEIPT_STUDIO_URL")

    if route == "direct":
        return Brain("direct", _norm(settings.llm_base_url), model, "RECEIPT_LLM_BASE_URL")

    if route == "hermes":
        hermes = [
            b
            for b in candidate_brains(settings)
            if b.kind == "hermes" or b.source.startswith("hermes") or b.source == "RECEIPT_HERMES_BASE_URL"
        ]
        if settings.hermes_base_url:
            hermes.insert(
                0,
                Brain("hermes", _norm(settings.hermes_base_url), settings.hermes_model or model, "RECEIPT_HERMES_BASE_URL"),
            )
        if not hermes:
            hermes = brains_from_hermes_file(settings.hermes_config_path)
        if not hermes:
            raise RuntimeError(
                "no Hermes Gemma 4 12B provider — run scripts/register-hermes-gemma.py "
                "or set RECEIPT_HERMES_BASE_URL"
            )
        check = probe or probe_models
        for brain in hermes:
            ids = check(brain.base_url)
            if ids is None:
                continue
            if not ids or any(looks_like_gemma(i) for i in ids) or brain.model in ids:
                return brain
        return hermes[0]

    # auto
    studio = (settings.studio_url or "").rstrip("/")
    if studio:
        ok = (studio_probe or (lambda u: probe_studio(u)))(studio)
        if ok:
            return Brain("studio", studio, model, "RECEIPT_STUDIO_URL")

    check = probe or probe_models
    for brain in candidate_brains(settings):
        ids = check(brain.base_url)
        if not ids:
            continue
        if any(looks_like_gemma(i) for i in ids) or brain.model in ids or brain.kind == "hermes":
            if brain.model not in ids and ids:
                gemma_id = next((i for i in ids if looks_like_gemma(i)), ids[0])
                return Brain(brain.kind, brain.base_url, gemma_id, brain.source)
            return brain
        # OpenAI-compat that lists nothing useful but is up: keep configured model
        if brain.source == "RECEIPT_LLM_BASE_URL":
            return brain

    # Prefer configured URL even if the probe failed (server may be starting).
    return Brain("direct", _norm(settings.llm_base_url), model, "RECEIPT_LLM_BASE_URL")