249 lines
9.7 KiB
Python
249 lines
9.7 KiB
Python
"""OpenRouter chat-completions client (S2 seam).
|
||
|
||
OpenAI-compatible direct client — zero heavy deps (httpx only). Key / base URL
|
||
/ model all come from the environment; the client is injectable for tests
|
||
(transport or a stub subclass). No call-quota logic by design (CONTRACT C3:
|
||
no artificial limits on real calls; only ordinary robustness retries).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
from typing import Any, Callable
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger("hvr.llm")
|
||
|
||
DEFAULT_MODEL = "deepseek/deepseek-v4-flash-0731"
|
||
DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"
|
||
|
||
_JSON_FENCE = re.compile(r"^```(?:json)?\s*|\s*```$")
|
||
|
||
|
||
class LLMError(Exception):
|
||
"""User-readable upstream failure (HTTP/network/parse-after-retry)."""
|
||
|
||
|
||
def resolve_env(env: dict[str, str] | None = None) -> dict[str, str]:
|
||
"""Resolve runtime config. Key stays external (CONTRACT C2)."""
|
||
env = env or dict(os.environ)
|
||
api_key = env.get("OPENROUTER_API_KEY") or env.get("PRODREAM_BACKEND_OPENROUTER_API_KEY") or ""
|
||
base_url = (env.get("OPENROUTER_BASE_URL") or DEFAULT_BASE_URL).rstrip("/")
|
||
# 正向代理（含 Basic Auth），如后端 .env 的 PRODREAM_BACKEND_OPENROUTER_PROXY_URL
|
||
proxy_url = env.get("OPENROUTER_PROXY_URL") or env.get("PRODREAM_BACKEND_OPENROUTER_PROXY_URL") or ""
|
||
model = env.get("HVR_LLM_MODEL") or DEFAULT_MODEL
|
||
return {"api_key": api_key, "base_url": base_url, "proxy_url": proxy_url, "model": model}
|
||
|
||
|
||
def _join_content_parts(content: Any) -> str:
|
||
if isinstance(content, str):
|
||
return content
|
||
if isinstance(content, list):
|
||
parts: list[str] = []
|
||
for part in content:
|
||
if isinstance(part, str):
|
||
parts.append(part)
|
||
elif isinstance(part, dict):
|
||
parts.append(str(part.get("text") or ""))
|
||
return "".join(parts)
|
||
return ""
|
||
|
||
|
||
def message_text(data: dict[str, Any]) -> str:
|
||
"""Read assistant text from an OpenAI-compatible completion payload.
|
||
|
||
DeepSeek V4 / OpenRouter reasoning models often put thinking in
|
||
``reasoning`` / ``reasoning_content`` and leave ``content`` empty when
|
||
the shared output budget is exhausted.
|
||
"""
|
||
try:
|
||
choice = data["choices"][0]
|
||
message = choice.get("message") or {}
|
||
except (KeyError, IndexError, TypeError) as exc:
|
||
raise LLMError("大模型响应缺少内容字段") from exc
|
||
text = _join_content_parts(message.get("content"))
|
||
if text.strip():
|
||
return text
|
||
finish = str(choice.get("finish_reason") or "")
|
||
has_reasoning = bool(
|
||
message.get("reasoning")
|
||
or message.get("reasoning_content")
|
||
or message.get("reasoning_details")
|
||
)
|
||
extra: list[str] = []
|
||
if finish:
|
||
extra.append(f"finish_reason={finish}")
|
||
if has_reasoning:
|
||
extra.append("思考过程占满了输出额度")
|
||
suffix = "；".join(extra) if extra else "可能被截断或拒绝"
|
||
raise LLMError(f"模型未返回内容（{suffix}）")
|
||
|
||
|
||
def reasoning_config(disable_reasoning: bool = False) -> dict[str, Any]:
|
||
"""OpenRouter rejects payloads that set both ``effort`` and ``max_tokens``."""
|
||
if disable_reasoning:
|
||
return {"effort": "none", "exclude": True}
|
||
return {"max_tokens": 2048}
|
||
|
||
|
||
def extract_json(content: str | None) -> dict[str, Any]:
|
||
"""Parse model output as JSON. Strips ```json fences and trims to the
|
||
outermost balanced {...} region."""
|
||
if not content or not content.strip():
|
||
raise LLMError("模型未返回内容（可能被截断或拒绝）")
|
||
text = _JSON_FENCE.sub("", content.strip()).strip()
|
||
try:
|
||
parsed = json.loads(text)
|
||
if isinstance(parsed, dict):
|
||
return parsed
|
||
except json.JSONDecodeError:
|
||
pass
|
||
start = text.find("{")
|
||
end = text.rfind("}")
|
||
if start != -1 and end > start:
|
||
try:
|
||
parsed = json.loads(text[start : end + 1])
|
||
if isinstance(parsed, dict):
|
||
return parsed
|
||
except json.JSONDecodeError:
|
||
pass
|
||
raise LLMError("模型输出不是有效的 JSON（解析失败）")
|
||
|
||
|
||
def _readable_error(resp: httpx.Response) -> str:
|
||
"""从上游错误响应提取一行用户可读信息（优先 error.message），
|
||
不把完整 JSON body 甩给页面（PRD §11 失败兜底）。"""
|
||
try:
|
||
body = resp.json()
|
||
if isinstance(body, dict):
|
||
msg = (body.get("error") or {}).get("message", "") if isinstance(body.get("error"), dict) else ""
|
||
if msg:
|
||
return msg[:300]
|
||
except Exception:
|
||
pass
|
||
return resp.text[:200]
|
||
|
||
|
||
class LlmClient:
|
||
"""Thin OpenAI-compatible client. `transport`/`client` injectable for tests."""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str | None = None,
|
||
base_url: str | None = None,
|
||
model: str | None = None,
|
||
proxy: str | None = None,
|
||
timeout: float = 180.0,
|
||
max_tokens: int = 16000,
|
||
temperature: float = 0.3,
|
||
transport: httpx.BaseTransport | None = None,
|
||
) -> None:
|
||
cfg = resolve_env()
|
||
self.api_key = api_key or cfg["api_key"]
|
||
self.base_url = base_url or cfg["base_url"]
|
||
self.model = model or cfg["model"]
|
||
self.proxy = proxy if proxy is not None else cfg["proxy_url"]
|
||
self.timeout = timeout
|
||
self.max_tokens = max_tokens
|
||
self.temperature = temperature
|
||
if not self.api_key:
|
||
raise LLMError("缺少 OpenRouter Key：请通过 local_start.sh 启动（从 prodream_backend/.env 读取）")
|
||
self._transport = transport
|
||
|
||
# -- internal ----------------------------------------------------------
|
||
def _post(self, payload: dict[str, Any]) -> httpx.Response:
|
||
kwargs: dict[str, Any] = {"timeout": self.timeout}
|
||
if self._transport is not None:
|
||
kwargs["transport"] = self._transport
|
||
# trust_env=False + 显式 proxy：直连配置的 base_url，不继承本机 Surge 等系统代理；
|
||
# 区域代理（PRODREAM_BACKEND_OPENROUTER_PROXY_URL）作为正向代理传入
|
||
with httpx.Client(trust_env=False, proxy=self.proxy or None, **kwargs) as client:
|
||
resp = client.post(
|
||
f"{self.base_url}/chat/completions",
|
||
headers={
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
json=payload,
|
||
)
|
||
return resp
|
||
|
||
def _complete(
|
||
self,
|
||
system: str,
|
||
user: str,
|
||
json_reminder: bool = False,
|
||
disable_reasoning: bool = False,
|
||
) -> str:
|
||
messages = [{"role": "system", "content": system}]
|
||
if json_reminder:
|
||
messages.append(
|
||
{"role": "system", "content": "只输出符合要求的 JSON，不要输出任何其它文字、解释或代码块标记。"}
|
||
)
|
||
messages.append({"role": "user", "content": user})
|
||
payload: dict[str, Any] = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"temperature": self.temperature,
|
||
"max_tokens": self.max_tokens,
|
||
"response_format": {"type": "json_object"},
|
||
}
|
||
# DeepSeek V4 Flash 的思考 token 与正文共用输出额度；不限制时
|
||
# diagnose 这类长 JSON 经常只返回空 content（页面上就是 502）。
|
||
# OpenRouter：effort 与 max_tokens 只能二选一。
|
||
payload["reasoning"] = reasoning_config(disable_reasoning)
|
||
try:
|
||
resp = self._post(payload)
|
||
except httpx.HTTPError as exc: # network / timeout
|
||
raise LLMError(f"调用大模型失败（网络/超时）：{exc.__class__.__name__}") from exc
|
||
if resp.status_code != 200:
|
||
raise LLMError(f"调用大模型失败（HTTP {resp.status_code}）：{_readable_error(resp)}")
|
||
return message_text(resp.json())
|
||
|
||
# -- public ------------------------------------------------------------
|
||
def complete_json(
|
||
self,
|
||
system: str,
|
||
user: str,
|
||
validate: Callable[[dict[str, Any]], str | None] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""Structured completion with robustness retries (SEAMS default:
|
||
any LLMError (parse failure / empty content / transient upstream
|
||
failure) is retried once, the retry carrying a JSON-only reminder
|
||
and turning reasoning off so the answer is not crowded out).
|
||
No quota logic.
|
||
|
||
`validate` is an optional semantic check on the parsed JSON: return
|
||
None when acceptable, or a one-line Chinese explanation when not
|
||
(e.g. per-paragraph arrays must match the input paragraph count).
|
||
A failing check is treated exactly like a parse failure — logged,
|
||
retried once, only surfaced after both attempts fail."""
|
||
last_error: LLMError | None = None
|
||
for attempt in range(2):
|
||
try:
|
||
content = self._complete(
|
||
system,
|
||
user,
|
||
json_reminder=(attempt == 1),
|
||
disable_reasoning=(attempt == 1),
|
||
)
|
||
parsed = extract_json(content)
|
||
if validate is not None:
|
||
problem = validate(parsed)
|
||
if problem:
|
||
raise LLMError(problem)
|
||
if attempt == 1:
|
||
logger.warning("llm retry succeeded after first failure: %s", last_error)
|
||
return parsed
|
||
except LLMError as exc:
|
||
last_error = exc
|
||
if attempt == 0:
|
||
logger.warning("llm call failed, retrying once: %s", exc)
|
||
assert last_error is not None
|
||
logger.error("llm call failed after 2 attempts: %s", last_error)
|
||
raise last_error
|