105 lines
3.7 KiB
Python
105 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import urllib.error
|
|
import urllib.request
|
|
from collections.abc import Iterator
|
|
from typing import Any, Protocol
|
|
|
|
from backend.llm.streaming import TextAccumulator
|
|
|
|
|
|
class RuntimeProfile(Protocol):
|
|
api_key: str
|
|
base_url: str
|
|
model_identifier: str
|
|
|
|
|
|
class ProviderFailure(RuntimeError):
|
|
def __init__(self, code: str) -> None:
|
|
super().__init__(code)
|
|
self.code = code
|
|
|
|
|
|
class OpenAICompatibleClient:
|
|
def __init__(self, timeout_seconds: int = 90) -> None:
|
|
self._timeout = timeout_seconds
|
|
|
|
def stream(
|
|
self, profile: RuntimeProfile, messages: list[dict[str, str]]
|
|
) -> Iterator[str]:
|
|
payload = json.dumps(
|
|
{"model": profile.model_identifier, "messages": messages, "stream": True},
|
|
ensure_ascii=False,
|
|
).encode("utf-8")
|
|
request = urllib.request.Request(
|
|
f"{profile.base_url.rstrip('/')}/chat/completions",
|
|
data=payload,
|
|
headers={
|
|
"Accept": "text/event-stream",
|
|
"Authorization": f"Bearer {profile.api_key}",
|
|
"Content-Type": "application/json",
|
|
"User-Agent": "XiaobaiReview/2",
|
|
},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self._timeout) as response:
|
|
accumulator = TextAccumulator()
|
|
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:
|
|
payload = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
choices = payload.get("choices") if isinstance(payload, dict) else None
|
|
if not isinstance(choices, list) or not choices:
|
|
continue
|
|
choice = choices[0]
|
|
chunk = accumulator.feed(choice if isinstance(choice, dict) else {})
|
|
if chunk:
|
|
yield chunk
|
|
if not accumulator.text:
|
|
raise ProviderFailure("empty_response")
|
|
except ProviderFailure:
|
|
raise
|
|
except urllib.error.HTTPError as exc:
|
|
raise ProviderFailure(_http_error(exc)) from exc
|
|
except TimeoutError as exc:
|
|
raise ProviderFailure("timeout") from exc
|
|
except (urllib.error.URLError, ConnectionError, OSError) as exc:
|
|
raise ProviderFailure("network") from exc
|
|
|
|
|
|
def _http_error(error: urllib.error.HTTPError) -> str:
|
|
if error.code in {401, 403}:
|
|
return "authentication"
|
|
if error.code == 429:
|
|
detail = _error_detail(error)
|
|
return "capacity" if "capacity" in detail.casefold() else "rate_limited"
|
|
if error.code in {408, 504}:
|
|
return "timeout"
|
|
if error.code >= 500:
|
|
detail = _error_detail(error)
|
|
return "capacity" if "capacity" in detail.casefold() else "upstream"
|
|
return "request_rejected"
|
|
|
|
|
|
def _error_detail(error: urllib.error.HTTPError) -> str:
|
|
try:
|
|
payload: Any = json.loads(error.read().decode("utf-8", errors="replace"))
|
|
except (json.JSONDecodeError, OSError):
|
|
return ""
|
|
if not isinstance(payload, dict):
|
|
return ""
|
|
detail = payload.get("error") or payload.get("message") or ""
|
|
if isinstance(detail, dict):
|
|
detail = detail.get("message") or detail.get("code") or ""
|
|
return str(detail)[:500]
|