providers.py

"""Model registry, client construction, and the key policy for public hosting.

Every provider here speaks the OpenAI wire format, so one client class covers
all of them and only the base URL changes.

Key policy, in priority order per request:

  1. A key the visitor typed into the UI ("bring your own key"). Used for that
     single request and then dropped - never logged, never persisted, never
     written to disk.
  2. Otherwise the host's own key for that provider, if one is configured.
     On the public deployment that is deliberately a free-tier provider, and
     it is rate limited per session.

The demo defaults to Groq's free tier serving gpt-oss-120b - which is the model
that actually won this project's benchmark, so the zero-cost path is also the
one that shows the best result.
"""

from __future__ import annotations

import os
import time
from dataclasses import dataclass, field

from openai import OpenAI

# --------------------------------------------------------------------------
# registry
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class Provider:
    key_env: str          # environment variable holding the key
    base_url: str | None  # None means the OpenAI default endpoint
    label: str
    signup: str = ""
    needs_key: bool = True


PROVIDERS: dict[str, Provider] = {
    "openai": Provider("OPENAI_API_KEY", None, "OpenAI",
                       "https://platform.openai.com/api-keys"),
    "anthropic": Provider("ANTHROPIC_API_KEY", "https://api.anthropic.com/v1/", "Anthropic",
                          "https://console.anthropic.com/settings/keys"),
    "google": Provider("GOOGLE_API_KEY", "https://generativelanguage.googleapis.com/v1beta/openai/", "Google",
                       "https://aistudio.google.com/apikey"),
    "grok": Provider("GROK_API_KEY", "https://api.x.ai/v1", "xAI",
                     "https://console.x.ai"),
    "groq": Provider("GROQ_API_KEY", "https://api.groq.com/openai/v1", "Groq (free tier)",
                     "https://console.groq.com/keys"),
    "openrouter": Provider("OPENROUTER_API_KEY", "https://openrouter.ai/api/v1", "OpenRouter",
                           "https://openrouter.ai/keys"),
    "ollama": Provider("", os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1"),
                       "Ollama (local)", "https://ollama.com", needs_key=False),
}


DEMO = "demo"      # cheap or free - the host's key may serve these
BYOK = "byok"      # expensive - the visitor must supply their own key
LOCAL = "local"    # runs on the machine hosting the app, costs nothing


@dataclass(frozen=True)
class Model:
    id: str
    provider: str
    display: str
    note: str = ""
    reasoning: bool = False  # accepts the reasoning_effort parameter
    tier: str = BYOK         # default to the safe side: never spend the host's money


# Tiering is based on measured results, not list price alone. On the max-subarray
# benchmark, gpt-5-nano produced a correct translation running 62x faster than
# Python; GPT-5 with high reasoning effort scored 63x. The difference is noise,
# and the cost difference is not - so the small model is the demo default and
# the frontier models are opt-in with the visitor's own key.
MODELS: list[Model] = [
    # -- demo tier: safe to serve with the host's key --------------------
    Model("openai/gpt-oss-120b", "groq", "GPT-OSS 120B",
          "free tier - fastest result in this benchmark", tier=DEMO),
    Model("gpt-5-nano", "openai", "GPT-5 Nano",
          "very low cost - scored 62x on the benchmark", reasoning=True, tier=DEMO),
    Model("gpt-5-mini", "openai", "GPT-5 Mini",
          "low cost - scored 63x", reasoning=True, tier=DEMO),

    # -- frontier: visitor's own key only --------------------------------
    Model("gpt-5", "openai", "GPT-5", "frontier", reasoning=True, tier=BYOK),
    Model("claude-sonnet-5", "anthropic", "Claude Sonnet 5", "frontier", tier=BYOK),
    Model("grok-4", "grok", "Grok 4", "frontier", tier=BYOK),
    Model("qwen/qwen3-coder-30b-a3b-instruct", "openrouter", "Qwen3 Coder 30B",
          "open weights", tier=BYOK),
    # Flash models failed to compile on both benchmarks tried - they emit Python
    # format strings such as {:.6f}, which is not valid Rust. Kept for comparison
    # because that failure is itself an interesting result, but not defaulted to.
    Model("gemini-3.5-flash", "google", "Gemini 3.5 Flash",
          "fast, but failed to compile in testing", tier=BYOK),

    # -- local via Ollama -------------------------------------------------
    Model("gpt-oss:20b", "ollama", "GPT-OSS 20B", "runs locally", tier=LOCAL),
    Model("qwen2.5-coder", "ollama", "Qwen2.5 Coder", "runs locally", tier=LOCAL),
    Model("deepseek-coder-v2", "ollama", "DeepSeek Coder V2", "runs locally", tier=LOCAL),
]

BY_ID = {m.id: m for m in MODELS}

DEMO_MODEL = os.getenv("DEMO_MODEL", "openai/gpt-oss-120b")


# --------------------------------------------------------------------------
# availability
# --------------------------------------------------------------------------


_ollama_cache: tuple[float, bool] | None = None


def ollama_reachable() -> bool:
    """Whether a local Ollama server is actually answering.

    Ollama needs no key, so a key check would report its models as ready even
    with nothing running - and the visitor would get a connection error after
    waiting. Probing the socket keeps the availability badges honest. Cached
    briefly so it does not run on every dropdown render.
    """
    global _ollama_cache
    now = time.time()
    if _ollama_cache and now - _ollama_cache[0] < 30:
        return _ollama_cache[1]

    import urllib.error
    import urllib.request

    base = PROVIDERS["ollama"].base_url or ""
    url = base.rstrip("/").removesuffix("/v1") + "/api/tags"
    try:
        with urllib.request.urlopen(url, timeout=1.5) as response:
            alive = response.status == 200
    except Exception:
        alive = False

    _ollama_cache = (now, alive)
    return alive


def host_key_for(provider: str) -> str | None:
    """The host's configured key for a provider, if any."""
    spec = PROVIDERS[provider]
    if not spec.needs_key:
        return "not-needed" if ollama_reachable() else None
    value = os.getenv(spec.key_env, "").strip()
    return value or None


def may_use_host_key(model: Model) -> bool:
    """Whether this model is allowed to spend the owner's credits.

    Frontier models never are - not on the public demo and not locally. The
    measured benefit over gpt-5-nano on this task is nil (62x vs 63x), so
    charging frontier prices for it is pure waste, and an accidental click
    while developing costs just as much as a visitor's.
    """
    if model.tier == LOCAL:
        return True          # costs nothing, runs on this machine
    return model.tier == DEMO


def visible_models() -> list[Model]:
    """Models worth showing on this instance.

    Local models are hidden unless Ollama is actually answering. The probe runs
    on the *server*, so a hosted container will never find one - and listing a
    permanently broken option is worse than not listing it.
    """
    models = [m for m in MODELS if m.tier != LOCAL]
    if not PUBLIC and ollama_reachable():
        models += [m for m in MODELS if m.tier == LOCAL]
    return models


def is_available(model_id: str) -> bool:
    """True when the host can serve this model with no key from the visitor."""
    model = BY_ID.get(model_id)
    if not model or not may_use_host_key(model):
        return False
    return host_key_for(model.provider) is not None


def needs_own_key(model_id: str) -> bool:
    """True when this model requires the visitor to supply a key."""
    model = BY_ID.get(model_id)
    if not model or model.tier == LOCAL:
        return False
    return not is_available(model_id)


def choice_label(model: Model) -> str:
    """Dropdown label. The cost warning goes first, where it cannot be missed.

    Trailing markers get truncated in a narrow select element - exactly the
    situation where someone is about to spend money without meaning to.
    """
    note = f" - {model.note}" if model.note else ""

    if is_available(model.id) or model.tier == LOCAL:
        return f"{model.display}{note}"
    if model.tier == DEMO and not PUBLIC:
        # Local dev: the owner can fix this, and naming the variable is the
        # fastest route. Suppressed in public mode - a visitor cannot set an
        # environment variable on someone else's server, so it would only
        # confuse them. There, the honest instruction is the same as frontier.
        return f"(set {PROVIDERS[model.provider].key_env})  {model.display}{note}"
    return f"(need your own key)  {model.display}{note}"


def dropdown_choices() -> list[tuple[str, str]]:
    """(label, value) pairs, ready models first."""
    models = visible_models()
    ordered = sorted(models, key=lambda m: (not is_available(m.id), models.index(m)))
    return [(choice_label(m), m.id) for m in ordered]


def default_model() -> str:
    """Prefer the configured demo model, then any ready model, then anything."""
    models = visible_models()
    if is_available(DEMO_MODEL):
        return DEMO_MODEL
    for m in models:
        if m.tier == DEMO and is_available(m.id):
            return m.id
    for m in models:
        if is_available(m.id):
            return m.id
    return models[0].id


def availability_summary() -> str:
    """Short markdown line describing what the running instance can serve."""
    models = visible_models()
    ready = [m for m in models if is_available(m.id)]
    if not ready:
        return (
            "**No models ready.** Add a provider key to `.env` "
            "(see `.env.example`), or paste your own key below to try it now."
        )
    names = ", ".join(m.display for m in ready[:3])
    extra = f" +{len(ready) - 3} more" if len(ready) > 3 else ""
    tail = ""
    if PUBLIC:
        tail = "  Frontier models need your own key - they are too costly to give away."
    return f"**Ready now:** {names}{extra}.{tail}"


# --------------------------------------------------------------------------
# rate limiting
# --------------------------------------------------------------------------


@dataclass
class RateLimiter:
    """Per-session cap on requests that spend the *host's* key.

    Requests carrying a visitor's own key are never counted - they cost the
    host nothing, so throttling them would only punish the considerate.

    In-memory and per-process, which is the right size for a single-container
    demo. A multi-replica deployment would need shared state.
    """

    limit: int = field(default_factory=lambda: int(os.getenv("DEMO_RATE_LIMIT", "5")))
    window_seconds: int = 3600
    _hits: dict[str, list[float]] = field(default_factory=dict)

    def check(self, session_id: str) -> tuple[bool, str]:
        if self.limit <= 0:
            return True, ""
        now = time.time()
        recent = [t for t in self._hits.get(session_id, []) if now - t < self.window_seconds]
        self._hits[session_id] = recent
        if len(recent) >= self.limit:
            wait = int((self.window_seconds - (now - recent[0])) / 60) + 1
            return False, (
                f"Demo limit reached ({self.limit} free conversions per hour).\n\n"
                f"Try again in ~{wait} minutes, or open **Use your own API key** "
                f"below for unlimited use with your own account."
            )
        return True, ""

    def record(self, session_id: str) -> None:
        self._hits.setdefault(session_id, []).append(time.time())

    def remaining(self, session_id: str) -> int:
        if self.limit <= 0:
            return -1
        now = time.time()
        recent = [t for t in self._hits.get(session_id, []) if now - t < self.window_seconds]
        return max(0, self.limit - len(recent))


limiter = RateLimiter()
PUBLIC = os.getenv("PUBLIC_DEPLOYMENT", "0") == "1"


# --------------------------------------------------------------------------
# client construction
# --------------------------------------------------------------------------


class NoKeyError(RuntimeError):
    """Raised when neither the visitor nor the host can supply a usable key."""


def build_client(model_id: str, user_key: str | None) -> tuple[OpenAI, bool]:
    """Return (client, used_host_key).

    `user_key` is used for this call only. It is never stored, logged, or
    written anywhere - it lives in this frame and is released with it.
    """
    model = BY_ID.get(model_id)
    if model is None:
        raise NoKeyError(f"Unknown model: {model_id}")

    spec = PROVIDERS[model.provider]
    user_key = (user_key or "").strip()

    if user_key:
        key, used_host = user_key, False
    elif not may_use_host_key(model):
        # Frontier tier. Refuse before the host key is even read, rather than
        # reading it and checking a limit afterwards.
        raise NoKeyError(
            f"**{model.display}** is a frontier model and never runs on this "
            f"instance's key - a few clicks is real money.\n\n"
            f"Open **Use your own API key** below and paste a {spec.label} key "
            f"({spec.signup}) to use it.\n\n"
            f"Or pick a model without the warning: on this benchmark GPT-5 Nano "
            f"scored 62x against GPT-5's 63x, so you give up essentially nothing."
        )
    else:
        host = host_key_for(model.provider)
        if host is None:
            if not spec.needs_key:
                raise NoKeyError(
                    f"**{model.display}** runs through a local Ollama server, and "
                    f"none is responding at {spec.base_url}.\n\n"
                    f"Start it with `ollama serve` and `ollama pull {model.id}`, "
                    f"or choose a hosted model instead. Pasting an API key will "
                    f"not help with a local model."
                )
            raise NoKeyError(
                f"**{model.display}** needs a {spec.label} key, and this instance "
                f"has none configured.\n\n"
                f"Open **Use your own API key** below and paste a {spec.label} key "
                f"({spec.signup}), or pick a model marked as ready."
            )
        key, used_host = host, True

    if not spec.needs_key:
        return OpenAI(api_key="ollama", base_url=spec.base_url), used_host
    if spec.base_url:
        return OpenAI(api_key=key, base_url=spec.base_url), used_host
    return OpenAI(api_key=key), used_host


def scrub(text: str) -> str:
    """Strip anything key-shaped out of text before it reaches the UI.

    Provider SDKs sometimes echo the offending credential back inside an error
    message. Since those messages get rendered into the page, they pass
    through here first.
    """
    import re

    patterns = [
        r"sk-[A-Za-z0-9_\-]{12,}",
        r"gsk_[A-Za-z0-9_\-]{12,}",
        r"xai-[A-Za-z0-9_\-]{12,}",
        r"AIza[A-Za-z0-9_\-]{20,}",
        r"sk-or-v1-[A-Za-z0-9_\-]{12,}",
        r"pplx-[A-Za-z0-9_\-]{12,}",
    ]
    for pattern in patterns:
        text = re.sub(pattern, "[redacted]", text)
    return text