116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from typing import Protocol, Sequence
|
||
|
||
import httpx
|
||
|
||
from .exceptions import (
|
||
LLMNetworkError,
|
||
LLMNotConfiguredError,
|
||
LLMResponseError,
|
||
LLMTimeoutError,
|
||
)
|
||
from .types import ChatMessage, TokenUsage
|
||
|
||
|
||
class LLMClient(Protocol):
|
||
"""LLM 调用适配器（可注入替换为 Fake）。"""
|
||
|
||
def chat(
|
||
self,
|
||
*,
|
||
model: str,
|
||
messages: list[ChatMessage],
|
||
temperature: float,
|
||
max_tokens: int,
|
||
) -> tuple[str, TokenUsage]: ...
|
||
|
||
|
||
class HttpLLMClient:
|
||
"""OpenAI Chat Completions 兼容的 httpx 实现；支持重试（指数退避）。"""
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
base_url: str,
|
||
api_key: str,
|
||
timeout_sec: float = 60.0,
|
||
retry_backoff: Sequence[float] = (1.0, 3.0, 7.0),
|
||
transport: httpx.BaseTransport | None = None,
|
||
) -> None:
|
||
if not api_key:
|
||
raise LLMNotConfiguredError("LLM API key 未配置（DEEPSEEK_API_KEY / LLM_BASE_URL）")
|
||
self._base_url = base_url.rstrip("/")
|
||
self._api_key = api_key
|
||
self._timeout_sec = timeout_sec
|
||
self._retry_backoff = retry_backoff
|
||
self._client = httpx.Client(timeout=timeout_sec, transport=transport)
|
||
|
||
def __enter__(self) -> HttpLLMClient:
|
||
"""支持 with 块：退出时自动关闭底层连接。"""
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
||
self._client.close()
|
||
|
||
def chat(
|
||
self,
|
||
*,
|
||
model: str,
|
||
messages: list[ChatMessage],
|
||
temperature: float,
|
||
max_tokens: int,
|
||
) -> tuple[str, TokenUsage]:
|
||
url = f"{self._base_url}/v1/chat/completions"
|
||
payload = {
|
||
"model": model,
|
||
"messages": [{"role": m.role, "content": m.content} for m in messages],
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
}
|
||
headers = {
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
|
||
attempts = 1 + len(self._retry_backoff)
|
||
last_error: Exception | None = None
|
||
for attempt in range(attempts):
|
||
if attempt > 0:
|
||
time.sleep(self._retry_backoff[attempt - 1])
|
||
try:
|
||
resp = self._client.post(url, json=payload, headers=headers)
|
||
except httpx.TimeoutException as exc:
|
||
last_error = exc
|
||
continue
|
||
except httpx.HTTPError as exc:
|
||
last_error = exc
|
||
continue
|
||
|
||
if resp.status_code >= 500:
|
||
last_error = LLMNetworkError(f"LLM 5xx: {resp.status_code}")
|
||
continue
|
||
if resp.status_code >= 400:
|
||
raise LLMNetworkError(f"LLM HTTP {resp.status_code}: {resp.text[:200]}")
|
||
if not (200 <= resp.status_code < 300):
|
||
# 3xx 重定向不自动跟随，不得误判为成功
|
||
raise LLMNetworkError(f"LLM HTTP {resp.status_code}: {resp.text[:200]}")
|
||
|
||
try:
|
||
data = resp.json()
|
||
content = data["choices"][0]["message"]["content"]
|
||
except (json.JSONDecodeError, KeyError, IndexError, TypeError) as exc:
|
||
# 2xx 但响应结构损坏（非 JSON / 缺字段）→ 结构化错误，不重试
|
||
raise LLMResponseError(f"LLM 响应结构损坏: {exc}") from exc
|
||
usage_raw = data.get("usage", {})
|
||
usage = TokenUsage(
|
||
input_tokens=usage_raw.get("prompt_tokens", 0),
|
||
output_tokens=usage_raw.get("completion_tokens", 0),
|
||
)
|
||
return content, usage
|
||
|
||
if isinstance(last_error, httpx.TimeoutException):
|
||
raise LLMTimeoutError(f"LLM 超时（{self._timeout_sec}s）") from last_error
|
||
raise LLMNetworkError(f"LLM 调用失败（重试耗尽）: {last_error}") from last_error |