models.py

"""Pollinations model catalog — fetched once per session, filtered by modality.

The hardcoded model list approach was brittle: aliases came and went, and
several entries were quietly stale (`claude-haiku-4.5`, `gemini-2.5-pro`).
This module talks to `/v1/models` at startup, caches the result for the
session, and exposes two filtered views for the UI dropdowns.

Defaults to a small curated fallback list if Pollinations is unreachable so
the app still boots offline.
"""
from __future__ import annotations

import os
from typing import Any

import requests

MODELS_URL = "https://gen.pollinations.ai/v1/models"

# Curated fallback — small, safe set. Used only if the catalog fetch fails.
# Same IDs as the previous hardcoded list, minus stale entries codex flagged.
_FALLBACK_TEXT = [
    "claude-fast", "claude", "claude-opus-4.7", "claude-large",
    "openai-fast", "openai", "openai-large",
    "gpt-5.4-mini", "gpt-5.4",
    "deepseek", "deepseek-pro",
    "grok", "grok-large",
    "qwen-large", "qwen-coder",
    "gemma", "step-flash", "step-3.5-flash",
]
_FALLBACK_AUDIO = ["openai-audio", "openai-audio-large", "gemini", "gemini-3-flash"]

# Models that are audio-input but only do transcription (whisper, scribe,
# universal-*) — useless for our brief-style narrative output. Exclude them
# from the C dropdown so the user doesn't pick a transcription-only model and
# get a flat dump of lyrics back instead of a structured brief.
_TRANSCRIPTION_ONLY = {"whisper", "scribe", "universal-2", "universal-3-pro"}

# Models marked as gemini-style audio-input but designed for live realtime
# session APIs, not single-shot chat completion. Skip in our use case.
_REALTIME_ONLY = {"gpt-realtime-2"}

# Audio-input gemini models support tool/code execution and routinely
# burn token budget on tool round-trips. We still let the user pick them
# (codex P5: tag as experimental) but mark them visibly. gemini-search-*
# pair audio with Google Search grounding — also experimental in our use.
_TOOLS_RISKY_AUDIO = {"gemini", "gemini-3-flash", "gemini-flash-lite-3.1",
                      "gemini-large", "gemini-search-fast", "gemini-search-large"}

# Preferred order for the audio dropdown — openai-audio family first because
# they're pure listen-and-answer with no tool loop. Then experimentals.
_AUDIO_PREFERRED = ["openai-audio", "openai-audio-large"]


# ---------------------------------------------------------------------------
# Catalog fetch (cached)
# ---------------------------------------------------------------------------

_cache: list[dict[str, Any]] | None = None


def _auth_header() -> dict[str, str]:
    for env in ("POLLINATIONS_API_KEY", "POLLINATIONS_TOKEN"):
        v = (os.environ.get(env) or "").strip()
        if v:
            return {"Authorization": f"Bearer {v}"}
    try:
        from wallet import stored_key
        k = (stored_key() or "").strip()
        if k:
            return {"Authorization": f"Bearer {k}"}
    except Exception:
        pass
    return {}


def fetch_catalog(force: bool = False, timeout: float = 4.0) -> list[dict[str, Any]]:
    """Return the cached /v1/models payload (data array). Fetches once per
    session unless force=True. Returns [] on failure — callers must handle."""
    global _cache
    if _cache is not None and not force:
        return _cache
    try:
        r = requests.get(MODELS_URL, headers=_auth_header(), timeout=timeout)
        r.raise_for_status()
        _cache = (r.json() or {}).get("data") or []
    except Exception:
        _cache = []
    return _cache


def text_models() -> list[str]:
    """Models accepting text input and producing text output, suitable for
    the measured-brief A/B columns. Excludes audio-input variants (those live
    in audio_models()) and transcription-only models."""
    catalog = fetch_catalog()
    if not catalog:
        return list(_FALLBACK_TEXT)
    ids: list[str] = []
    for m in catalog:
        mid = m.get("id") or ""
        if not mid:
            continue
        inp = m.get("input_modalities") or []
        out = m.get("output_modalities") or []
        # text->text models, no audio input. Vision-capable models with
        # image input are fine (they just won't be sent images here).
        if "text" in out and "audio" not in inp:
            # Skip search/embedding/coder oddities by endpoint check
            endpoints = m.get("supported_endpoints") or []
            if "/v1/chat/completions" not in endpoints:
                continue
            ids.append(mid)
    # Sort with curated favourites first (claude / openai / gemini-fast / etc.)
    return _sort_with_favourites(ids, _FALLBACK_TEXT)


def audio_models() -> list[tuple[str, bool]]:
    """Models that accept audio INPUT and produce text — for the audio-only
    C column. Returns list of (id, is_experimental) tuples; experimental
    models are gemini ones that use tool/code-execution and may eat the
    token budget without producing a prose answer."""
    catalog = fetch_catalog()
    if not catalog:
        return [(m, False) for m in _FALLBACK_AUDIO[:2]] + \
               [(m, True) for m in _FALLBACK_AUDIO[2:]]
    pairs: list[tuple[str, bool]] = []
    for m in catalog:
        mid = m.get("id") or ""
        if not mid or mid in _TRANSCRIPTION_ONLY or mid in _REALTIME_ONLY:
            continue
        inp = m.get("input_modalities") or []
        out = m.get("output_modalities") or []
        if "audio" not in inp or "text" not in out:
            continue
        endpoints = m.get("supported_endpoints") or []
        if "/v1/chat/completions" not in endpoints:
            continue
        is_experimental = mid in _TOOLS_RISKY_AUDIO
        pairs.append((mid, is_experimental))
    # Stable ordering — non-experimental first, with curated favourites at
    # the very top within their group.
    def _sort_key(p: tuple[str, bool]) -> tuple[int, int, str]:
        mid, exp = p
        pref_idx = _AUDIO_PREFERRED.index(mid) if mid in _AUDIO_PREFERRED else 999
        return (int(exp), pref_idx, mid)
    pairs.sort(key=_sort_key)
    return pairs


def audio_model_choices() -> list[tuple[str, str]]:
    """UI-friendly form: (display_label, value) pairs.

    No more `· experimental` suffix on gemini — the salvage-from-content-blocks
    path + tool_choice:none fallback in narrative.py mean gemini's code-exec
    behaviour no longer silently swallows the prose answer. Treating all
    audio-input models as first-class makes the Compare default cleaner.
    """
    return [(mid, mid) for mid, _exp in audio_models()]


def _sort_with_favourites(ids: list[str], favourites: list[str]) -> list[str]:
    """Stable sort: keep `favourites` order at the front, everything else
    alphabetical after. Misses in favourites are silently skipped."""
    seen = set(ids)
    head = [m for m in favourites if m in seen]
    tail = sorted(m for m in ids if m not in head)
    return head + tail