xiaobai-review/backend/llm/transport.py

from __future__ import annotations

import json
import time
import urllib.error
import urllib.request
from collections.abc import Iterator
from dataclasses import dataclass
from typing import Any

from .stream import OpenAIStreamAccumulator


class OpenAITransportError(RuntimeError):
    pass


class OpenAIHTTPError(OpenAITransportError):
    def __init__(self, code: int, detail: str = "") -> None:
        super().__init__(f"HTTP {code}")
        self.code = code
        self.detail = detail

    def describe(self, label: str) -> str:
        suffix = f"：{self.detail[:300]}" if self.detail else ""
        return f"{label}（HTTP {self.code}）{suffix}"


class OpenAIEmptyResponseError(OpenAITransportError):
    pass


@dataclass(frozen=True)
class OpenAIChatCompletion:
    content: Any
    latency_ms: int


def chat_completion(
    *,
    api_key: str,
    base_url: str,
    model: str,
    messages: list[dict[str, Any]],
    timeout: int,
    user_agent: str,
) -> OpenAIChatCompletion:
    request = _request(api_key, base_url, model, messages, user_agent, stream=False)
    started = time.perf_counter()
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            result = json.loads(response.read().decode("utf-8"))
        content = result["choices"][0]["message"]["content"]
    except urllib.error.HTTPError as exc:
        raise OpenAIHTTPError(exc.code, _http_error_detail(exc)) from exc
    except (
        urllib.error.URLError,
        TimeoutError,
        json.JSONDecodeError,
        KeyError,
        IndexError,
    ) as exc:
        raise OpenAITransportError(str(exc)) from exc
    return OpenAIChatCompletion(
        content=content,
        latency_ms=round((time.perf_counter() - started) * 1000),
    )


def list_models(
    *,
    api_key: str,
    base_url: str,
    timeout: int,
    user_agent: str,
) -> list[str]:
    request = urllib.request.Request(
        f"{base_url.rstrip('/')}/models",
        headers={
            "Authorization": f"Bearer {api_key}",
            "User-Agent": user_agent,
        },
        method="GET",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            payload = json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        raise OpenAIHTTPError(exc.code, _http_error_detail(exc)) from exc
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, OSError) as exc:
        raise OpenAITransportError(str(exc)) from exc
    items = payload.get("data") if isinstance(payload, dict) else payload
    models = []
    for item in items or []:
        name = str((item or {}).get("id") or "") if isinstance(item, dict) else str(item or "")
        if name and name not in models:
            models.append(name)
    return models


def stream_chat_completion(
    *,
    api_key: str,
    base_url: str,
    model: str,
    messages: list[dict[str, Any]],
    timeout: int,
    user_agent: str,
) -> Iterator[str]:
    request = _request(api_key, base_url, model, messages, user_agent, stream=True)
    yielded = False
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            accumulator = OpenAIStreamAccumulator()
            for raw_line in response:
                line = raw_line.decode("utf-8", errors="replace").strip()
                if not line or line.startswith(":"):
                    continue
                if line.startswith("data:"):
                    line = line[5:].strip()
                if line == "[DONE]":
                    break
                try:
                    result = json.loads(line)
                except json.JSONDecodeError:
                    continue
                choices = result.get("choices") or []
                if not choices:
                    continue
                content = accumulator.feed(choices[0] or {})
                if content:
                    yielded = True
                    yield str(content)
    except urllib.error.HTTPError as exc:
        raise OpenAIHTTPError(exc.code, _http_error_detail(exc)) from exc
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        raise OpenAITransportError(str(exc)) from exc
    if not yielded:
        raise OpenAIEmptyResponseError("empty response")


def _request(
    api_key: str,
    base_url: str,
    model: str,
    messages: list[dict[str, Any]],
    user_agent: str,
    *,
    stream: bool,
) -> urllib.request.Request:
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
        "User-Agent": user_agent,
    }
    if stream:
        headers["Accept"] = "text/event-stream"
    return urllib.request.Request(
        f"{base_url.rstrip('/')}/chat/completions",
        data=json.dumps(
            {"model": model, "messages": messages, "stream": stream},
            ensure_ascii=False,
        ).encode("utf-8"),
        headers=headers,
        method="POST",
    )


def _http_error_detail(exc: urllib.error.HTTPError) -> str:
    try:
        payload = json.loads(exc.read().decode("utf-8", errors="replace"))
        error = payload.get("error")
        if isinstance(error, dict):
            return str(error.get("message") or error.get("code") or "")
        if error:
            return str(error)
        return str(payload.get("message") or "")
    except (json.JSONDecodeError, OSError):
        return ""