provider.py

"""OpenRouter async client with multimodal (image) support."""
from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass
from typing import Optional

import httpx

DEFAULT_TIMEOUT = 120.0
BASE_URL = "https://openrouter.ai/api/v1"


@dataclass
class ChatResult:
    model: str
    ok: bool
    content: str = ""
    latency_s: float = 0.0
    error: str = ""
    usage: Optional[dict] = None  # {prompt_tokens, completion_tokens, total_tokens, cost?}


def _headers(api_key: str) -> dict:
    return {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://htr-vlm-app.local",
        "X-Title": "HTR VLM + MoE Post-correction",
    }


def text_part(text: str) -> dict:
    return {"type": "text", "text": text}


def image_part(image_b64: str, mime: str = "image/jpeg") -> dict:
    return {
        "type": "image_url",
        "image_url": {"url": f"data:{mime};base64,{image_b64}"},
    }


class LLMClient:
    def __init__(self, api_key: str, timeout: float = DEFAULT_TIMEOUT):
        self.api_key = api_key
        self.timeout = timeout
        self.endpoint = BASE_URL + "/chat/completions"
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self):
        self._client = httpx.AsyncClient(
            timeout=self.timeout,
            limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
        )
        return self

    async def __aexit__(self, exc_type, exc, tb):
        if self._client:
            await self._client.aclose()
            self._client = None

    async def chat(
        self,
        *,
        model: str,
        system: str,
        user_parts: list[dict],
        temperature: float = 0.0,
        max_tokens: int = 4096,
        force_json: bool = True,
        max_retries: int = 2,
    ) -> ChatResult:
        if not self.api_key:
            return ChatResult(model=model, ok=False, error="Missing OpenRouter API key.")
        if self._client is None:
            raise RuntimeError("LLMClient must be used inside `async with` block")

        msgs = [
            {"role": "system", "content": system},
            {"role": "user", "content": user_parts},
        ]
        payload: dict = {
            "model": model,
            "messages": msgs,
            "temperature": temperature,
            "max_tokens": max_tokens,
            # OpenRouter usage-accounting: include token + cost details in response
            "usage": {"include": True},
        }
        if force_json:
            payload["response_format"] = {"type": "json_object"}

        start = time.time()
        last_error = ""
        backoff = 1.0

        for attempt in range(max_retries + 1):
            try:
                resp = await self._client.post(self.endpoint, headers=_headers(self.api_key), json=payload)
                if resp.status_code == 400 and force_json:
                    # Some routes reject response_format — strip and retry once.
                    payload_no_fmt = {k: v for k, v in payload.items() if k != "response_format"}
                    resp = await self._client.post(self.endpoint, headers=_headers(self.api_key), json=payload_no_fmt)
                if resp.status_code >= 400:
                    last_error = f"HTTP {resp.status_code}: {resp.text[:300]}"
                else:
                    data = resp.json()
                    content = data.get("choices", [{}])[0].get("message", {}).get("content") or ""
                    if not content and isinstance(data.get("choices", [{}])[0].get("message", {}).get("reasoning"), str):
                        content = data["choices"][0]["message"]["reasoning"]
                    usage = data.get("usage") or None
                    if isinstance(usage, dict):
                        usage = {
                            "prompt_tokens": usage.get("prompt_tokens"),
                            "completion_tokens": usage.get("completion_tokens"),
                            "total_tokens": usage.get("total_tokens"),
                            "cost": usage.get("cost"),
                        }
                    return ChatResult(
                        model=model,
                        ok=True,
                        content=content,
                        latency_s=time.time() - start,
                        usage=usage,
                    )
            except (httpx.TimeoutException, httpx.TransportError) as e:
                last_error = f"{type(e).__name__}: {e}"
            except Exception as e:  # noqa: BLE001
                last_error = f"{type(e).__name__}: {e}"

            if attempt < max_retries:
                await asyncio.sleep(backoff)
                backoff *= 2

        return ChatResult(
            model=model,
            ok=False,
            error=last_error or "unknown error",
            latency_s=time.time() - start,
        )


def test_connection_sync(api_key: str, model: Optional[str] = None) -> tuple[bool, str]:
    """Quick blocking test for the 'Test' button in the API key modal."""
    model = model or "mistralai/mistral-large-2512"
    try:
        resp = httpx.post(
            BASE_URL + "/chat/completions",
            headers=_headers(api_key),
            json={
                "model": model,
                "messages": [{"role": "user", "content": "Reply with the single word: OK"}],
                "max_tokens": 5,
                "temperature": 0.0,
            },
            timeout=20.0,
        )
        if resp.status_code >= 400:
            return False, f"HTTP {resp.status_code} from OpenRouter: {resp.text[:200]}"
        content = resp.json()["choices"][0]["message"]["content"]
        return True, f"OpenRouter key accepted."
    except Exception as e:  # noqa: BLE001
        return False, f"OpenRouter request failed: {e}"