provider.py

"""Request-scoped BYOK gateway for CaseAnchor. The user's API key lives
in the request body, goes only to the chosen provider, and is never stored,
logged, or echoed.

Presets cover OpenAI-compatible providers plus Anthropic's native format;
Gemini is reached through Google's OpenAI-compatibility endpoint. `custom`
accepts any OpenAI-compatible base URL (relays/中转站, LM Studio, Ollama,
vLLM) — HTTPS required except for localhost. `mock` needs no key and keeps
the whole product usable offline.
"""
from __future__ import annotations

import json
import socket
import urllib.error
import urllib.request
from typing import Any, Dict, List, Optional


class ProviderError(RuntimeError):
    def __init__(self, code: str, message: str, hint: str = ""):
        super().__init__(message)
        self.code = code
        self.hint = hint

    def to_dict(self) -> Dict[str, str]:
        return {"code": self.code, "message": str(self), "hint": self.hint}


PRESETS: Dict[str, Dict[str, str]] = {
    "openai":   {"url": "https://api.openai.com/v1/chat/completions", "format": "openai",
                 "default_model": "gpt-4o-mini"},
    "anthropic": {"url": "https://api.anthropic.com/v1/messages", "format": "anthropic",
                  "default_model": "claude-haiku-4-5-20251001"},
    "gemini":   {"url": "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
                 "format": "openai", "default_model": "gemini-2.0-flash"},
    "deepseek": {"url": "https://api.deepseek.com/chat/completions", "format": "openai",
                 "default_model": "deepseek-chat"},
    "kimi":     {"url": "https://api.moonshot.cn/v1/chat/completions", "format": "openai",
                 "default_model": "moonshot-v1-8k"},
    "glm":      {"url": "https://open.bigmodel.cn/api/paas/v4/chat/completions", "format": "openai",
                 "default_model": "glm-4-flash"},
}

TIMEOUT_SECONDS = 90

# Anthropic model families that reject sampling params (temperature/top_p)
# with HTTP 400 — send no temperature to these.
_NO_SAMPLING_MARKERS = ("opus-4-7", "opus-4-8", "fable", "mythos")


def _accepts_sampling(model: str) -> bool:
    m = (model or "").lower()
    return not any(marker in m for marker in _NO_SAMPLING_MARKERS)


def _is_local(url: str) -> bool:
    return any(url.startswith(prefix) for prefix in
               ("http://127.0.0.1", "http://localhost", "http://[::1]"))


def resolve_config(config: Dict[str, Any]) -> Dict[str, Any]:
    """Validate shape and resolve endpoint/format. Raises ProviderError."""
    provider = str(config.get("provider") or "mock").strip().lower()
    if provider == "mock":
        return {"provider": "mock", "format": "mock", "url": "", "model": "mock",
                "api_key": ""}
    api_key = str(config.get("api_key") or "").strip()
    model = str(config.get("model") or "").strip()
    if not api_key:
        raise ProviderError("missing_key", "API key is empty.",
                            "Paste a key in AI settings, or switch to mock mode.")
    if provider == "custom":
        base_url = str(config.get("base_url") or "").strip().rstrip("/")
        if not base_url:
            raise ProviderError("missing_base_url", "Custom base URL is empty.",
                                "Example: https://my-relay.example.com/v1")
        if not base_url.startswith("https://") and not _is_local(base_url):
            raise ProviderError("insecure_base_url",
                                "Custom base URL must use HTTPS (plain HTTP is "
                                "allowed only for localhost).",
                                "Use https://... or http://127.0.0.1:PORT")
        if not model:
            raise ProviderError("missing_model", "Model name is required for a "
                                "custom endpoint.", "Example: gpt-4o-mini")
        url = base_url + ("/chat/completions" if not base_url.endswith("/chat/completions") else "")
        return {"provider": provider, "format": "openai", "url": url,
                "model": model, "api_key": api_key}
    preset = PRESETS.get(provider)
    if preset is None:
        raise ProviderError("unknown_provider", f"Unknown provider '{provider}'.",
                            "Pick a preset or use 'custom' with a base URL.")
    return {"provider": provider, "format": preset["format"], "url": preset["url"],
            "model": model or preset["default_model"], "api_key": api_key}


def _post(url: str, headers: Dict[str, str], body: Dict[str, Any],
          timeout: int) -> Dict[str, Any]:
    payload = json.dumps(body).encode("utf-8")
    request = urllib.request.Request(url, data=payload, method="POST")
    request.add_header("Content-Type", "application/json")
    for name, value in headers.items():
        request.add_header(name, value)
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8", "replace"))
    except urllib.error.HTTPError as exc:
        detail = ""
        try:
            raw = exc.read().decode("utf-8", "replace")[:400]
            detail = json.loads(raw).get("error", {}).get("message", "") or raw
        except Exception:
            pass
        if exc.code in (401, 403):
            raise ProviderError("invalid_key", "The provider rejected the API key.",
                                "Re-check the key, its permissions, and whether "
                                "it matches the selected provider.")
        if exc.code == 404:
            raise ProviderError("model_missing",
                                f"Endpoint or model not found (HTTP 404). {detail}".strip(),
                                "Check the model name and base URL.")
        if exc.code == 429:
            raise ProviderError("rate_limited",
                                "Rate limit or quota exhausted (HTTP 429).",
                                "Wait and retry, or check the account balance.")
        if exc.code >= 500:
            raise ProviderError("provider_down",
                                f"Provider server error (HTTP {exc.code}).",
                                "Retry later; the provider side is failing.")
        raise ProviderError("http_error", f"HTTP {exc.code}: {detail}".strip(), "")
    except socket.timeout:
        raise ProviderError("timeout", f"No response within {timeout}s.",
                            "Slow model or network; retry or reduce batch size.")
    except urllib.error.URLError as exc:
        raise ProviderError("network", f"Network failure: {exc.reason}",
                            "Check connectivity, VPN/proxy, and the base URL.")
    except json.JSONDecodeError:
        raise ProviderError("bad_response", "Provider returned non-JSON output.",
                            "The base URL may not be an OpenAI-compatible endpoint.")


def chat(config: Dict[str, Any], system: str, user: str,
         max_tokens: int = 1200, temperature: float = 0.2,
         timeout: int = TIMEOUT_SECONDS) -> Dict[str, Any]:
    """One chat completion. Returns {text, usage{prompt_tokens, completion_tokens}}."""
    resolved = resolve_config(config)
    if resolved["format"] == "mock":
        return {"text": "[mock] Deterministic offline mode: no model was called.",
                "usage": {"prompt_tokens": 0, "completion_tokens": 0},
                "model": "mock"}
    if resolved["format"] == "anthropic":
        body: Dict[str, Any] = {"model": resolved["model"],
                                "max_tokens": max_tokens, "system": system,
                                "messages": [{"role": "user", "content": user}]}
        if _accepts_sampling(resolved["model"]):
            body["temperature"] = temperature
        data = _post(resolved["url"],
                     {"x-api-key": resolved["api_key"],
                      "anthropic-version": "2023-06-01"},
                     body, timeout)
        blocks = data.get("content") or []
        text = "".join(b.get("text", "") for b in blocks
                       if isinstance(b, dict) and b.get("type") == "text")
        if not text.strip():
            raise ProviderError("empty_response", "Provider returned no text.",
                                "Retry; if it persists, try another model.")
        usage = data.get("usage") or {}
        return {"text": text, "model": resolved["model"],
                "usage": {"prompt_tokens": int(usage.get("input_tokens") or 0),
                          "completion_tokens": int(usage.get("output_tokens") or 0)}}
    data = _post(resolved["url"],
                 {"Authorization": f"Bearer {resolved['api_key']}"},
                 {"model": resolved["model"], "temperature": temperature,
                  "max_tokens": max_tokens,
                  "messages": [{"role": "system", "content": system},
                               {"role": "user", "content": user}]},
                 timeout)
    choices = data.get("choices") or []
    if not choices:
        detail = str(data.get("error", {}).get("message", ""))[:200]
        raise ProviderError("empty_response",
                            f"Provider returned no choices. {detail}".strip(),
                            "Check the model name; some relays need exact names.")
    text = ((choices[0].get("message") or {}).get("content") or "").strip()
    if not text:
        raise ProviderError("empty_response", "Provider returned empty content.",
                            "Retry; if it persists, try another model.")
    usage = data.get("usage") or {}
    return {"text": text, "model": resolved["model"],
            "usage": {"prompt_tokens": int(usage.get("prompt_tokens") or 0),
                      "completion_tokens": int(usage.get("completion_tokens") or 0)}}


def validate_key(config: Dict[str, Any]) -> Dict[str, Any]:
    """Cheap live check of provider + key + model. Never raises on provider
    errors — returns {ok, code, message, hint} for the UI."""
    try:
        resolved = resolve_config(config)
        if resolved["format"] == "mock":
            return {"ok": True, "message": "Mock mode is offline and needs no key.",
                    "model": "mock"}
        result = chat(config, "Reply with exactly: ok", "ping",
                      max_tokens=8, temperature=0.0, timeout=30)
        return {"ok": True, "message": f"Key accepted by {resolved['provider']} "
                                       f"({result['model']}).",
                "model": result["model"], "usage": result["usage"]}
    except ProviderError as exc:
        payload = exc.to_dict()
        payload["ok"] = False
        return payload