72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
"""Chat-Client für das OpenAI-kompatible API von LiteLLM."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Protocol
|
|
|
|
import httpx
|
|
|
|
from .problems import is_context_overflow_message
|
|
|
|
|
|
class LLMError(Exception):
|
|
"""Basisklasse; `infra=True` bedeutet: Endpunkt/Netz, nicht Modellqualität."""
|
|
|
|
infra = True
|
|
|
|
|
|
class EndpointError(LLMError):
|
|
pass
|
|
|
|
|
|
class ContextOverflowError(LLMError):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class ChatResult:
|
|
content: str
|
|
finish_reason: str | None = None
|
|
usage: dict[str, Any] = field(default_factory=dict)
|
|
model: str | None = None
|
|
|
|
|
|
class ChatClient(Protocol):
|
|
def chat(self, model: str, messages: list[dict[str, str]], *, max_tokens: int | None = None,
|
|
temperature: float | None = None) -> ChatResult: ...
|
|
|
|
|
|
class OpenAIChatClient:
|
|
def __init__(self, base_url: str, api_key: str = "", timeout: float = 120.0,
|
|
client: httpx.Client | None = None):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
|
self.client = client or httpx.Client(timeout=timeout)
|
|
|
|
def chat(self, model: str, messages: list[dict[str, str]], *, max_tokens: int | None = None,
|
|
temperature: float | None = None) -> ChatResult:
|
|
body: dict[str, Any] = {"model": model, "messages": messages}
|
|
if max_tokens is not None:
|
|
body["max_tokens"] = max_tokens
|
|
if temperature is not None:
|
|
body["temperature"] = temperature
|
|
try:
|
|
resp = self.client.post(f"{self.base_url}/v1/chat/completions", json=body, headers=self.headers)
|
|
except httpx.HTTPError as exc:
|
|
raise EndpointError(f"{type(exc).__name__}: {exc}") from exc
|
|
if resp.status_code == 400 and is_context_overflow_message(resp.text):
|
|
raise ContextOverflowError(resp.text[:500])
|
|
if resp.status_code == 429 or resp.status_code >= 500:
|
|
raise EndpointError(f"HTTP {resp.status_code}: {resp.text[:300]}")
|
|
if resp.status_code >= 400:
|
|
raise EndpointError(f"HTTP {resp.status_code}: {resp.text[:300]}")
|
|
data = resp.json()
|
|
choice = data["choices"][0]
|
|
return ChatResult(
|
|
content=choice["message"].get("content") or "",
|
|
finish_reason=choice.get("finish_reason"),
|
|
usage=data.get("usage") or {},
|
|
model=data.get("model"),
|
|
)
|