276 lines
9.4 KiB
Python
276 lines
9.4 KiB
Python
"""ZenMux API 客户端。
|
||
|
||
直接 HTTP 调用 zenmux 的 OpenAI 兼容端点。为 Phase 4 的 Python 化脚本服务
|
||
（translate.py / polish.py）。所有调用都走 /api/v1/chat/completions。
|
||
|
||
设计原则：
|
||
- 独立于 opencode，可直接在 CLI / CI / cron 运行
|
||
- 内置重试（指数退避）、限流、token 统计、结构化日志
|
||
- 失败快速可见：打印请求 ID，便于在 zenmux 后台对账
|
||
- 默认读 secrets.env 里的 ZENMUX_API_KEY
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
DEFAULT_BASE_URL = "https://zenmux.ai/api/v1"
|
||
DEFAULT_TIMEOUT = 300.0 # 翻译/润色单次调用可能 60s+，留够余量
|
||
MAX_RETRIES = 5
|
||
RETRYABLE_STATUSES = {408, 429, 500, 502, 503, 504, 520, 524}
|
||
|
||
|
||
@dataclass
|
||
class UsageStats:
|
||
"""聚合一次脚本运行的 token 消耗。"""
|
||
|
||
prompt_tokens: int = 0
|
||
completion_tokens: int = 0
|
||
cache_creation_tokens: int = 0
|
||
cache_read_tokens: int = 0
|
||
call_count: int = 0
|
||
failed_calls: int = 0
|
||
by_model: dict[str, dict[str, int]] = field(default_factory=dict)
|
||
|
||
def add(self, model: str, usage: dict[str, Any]) -> None:
|
||
p = usage.get("prompt_tokens", 0) or 0
|
||
c = usage.get("completion_tokens", 0) or 0
|
||
cc = (
|
||
usage.get("cache_creation_input_tokens", 0)
|
||
or usage.get("prompt_tokens_details", {}).get("cached_tokens", 0)
|
||
or 0
|
||
)
|
||
cr = (
|
||
usage.get("cache_read_input_tokens", 0)
|
||
or 0
|
||
)
|
||
self.prompt_tokens += p
|
||
self.completion_tokens += c
|
||
self.cache_creation_tokens += cc
|
||
self.cache_read_tokens += cr
|
||
self.call_count += 1
|
||
m = self.by_model.setdefault(
|
||
model,
|
||
{"prompt": 0, "completion": 0, "cache_creation": 0, "cache_read": 0, "calls": 0},
|
||
)
|
||
m["prompt"] += p
|
||
m["completion"] += c
|
||
m["cache_creation"] += cc
|
||
m["cache_read"] += cr
|
||
m["calls"] += 1
|
||
|
||
def summary(self) -> str:
|
||
lines = [
|
||
f"Total calls: {self.call_count} (failed: {self.failed_calls})",
|
||
f"Prompt tokens: {self.prompt_tokens:>12,}",
|
||
f"Completion tokens: {self.completion_tokens:>12,}",
|
||
f"Cache creation: {self.cache_creation_tokens:>12,}",
|
||
f"Cache read: {self.cache_read_tokens:>12,}",
|
||
]
|
||
for model, s in self.by_model.items():
|
||
lines.append(
|
||
f" [{model}] calls={s['calls']} in={s['prompt']:,} "
|
||
f"out={s['completion']:,} cache_w={s['cache_creation']:,} "
|
||
f"cache_r={s['cache_read']:,}"
|
||
)
|
||
return "\n".join(lines)
|
||
|
||
|
||
class ZenMuxError(RuntimeError):
|
||
pass
|
||
|
||
|
||
class ZenMuxClient:
|
||
"""ZenMux 轻量客户端。
|
||
|
||
只暴露一个方法 `chat_complete()`，屏蔽 httpx 细节。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str | None = None,
|
||
base_url: str | None = None,
|
||
timeout: float = DEFAULT_TIMEOUT,
|
||
log_file: Path | None = None,
|
||
) -> None:
|
||
self.api_key = api_key or os.environ.get("ZENMUX_API_KEY")
|
||
if not self.api_key:
|
||
raise ZenMuxError(
|
||
"ZENMUX_API_KEY not set. Source secrets.env or pass api_key explicitly."
|
||
)
|
||
self.base_url = (base_url or os.environ.get("ZENMUX_BASE_URL") or DEFAULT_BASE_URL).rstrip("/")
|
||
self.timeout = timeout
|
||
self.log_file = log_file
|
||
self.usage = UsageStats()
|
||
self._client = httpx.Client(timeout=timeout)
|
||
self._log_lock = threading.Lock()
|
||
self._usage_lock = threading.Lock()
|
||
|
||
def close(self) -> None:
|
||
self._client.close()
|
||
|
||
def __enter__(self) -> "ZenMuxClient":
|
||
return self
|
||
|
||
def __exit__(self, *_args: Any) -> None:
|
||
self.close()
|
||
|
||
def _log(self, payload: dict[str, Any]) -> None:
|
||
if not self.log_file:
|
||
return
|
||
self.log_file.parent.mkdir(parents=True, exist_ok=True)
|
||
with self._log_lock:
|
||
with self.log_file.open("a", encoding="utf-8") as f:
|
||
f.write(json.dumps(payload, ensure_ascii=False) + "\n")
|
||
|
||
def chat_complete(
|
||
self,
|
||
model: str,
|
||
system: str,
|
||
user: str,
|
||
*,
|
||
temperature: float = 0.3,
|
||
max_tokens: int = 16000,
|
||
extra_messages: list[dict[str, str]] | None = None,
|
||
tag: str = "",
|
||
) -> str:
|
||
"""一次非流式对话补全。
|
||
|
||
Args:
|
||
model: 完整 model id，例如 `anthropic/claude-sonnet-4.6`（zenmux slug 不带 `zenmux/` 前缀，因为 baseURL 已经定位到 zenmux）。
|
||
system: system prompt
|
||
user: user message
|
||
temperature, max_tokens: 常规参数
|
||
extra_messages: 插在 system 之后、user 之前的额外消息（few-shot 等）
|
||
tag: 给这次调用打标签，便于日志里识别（如 "translate:ch03"）
|
||
|
||
Returns:
|
||
assistant 的纯文本内容。如失败抛 ZenMuxError。
|
||
"""
|
||
messages: list[dict[str, str]] = [{"role": "system", "content": system}]
|
||
if extra_messages:
|
||
messages.extend(extra_messages)
|
||
messages.append({"role": "user", "content": user})
|
||
|
||
body: dict[str, Any] = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
}
|
||
headers = {
|
||
"Authorization": f"Bearer {self.api_key}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
url = f"{self.base_url}/chat/completions"
|
||
|
||
last_error: str = ""
|
||
for attempt in range(MAX_RETRIES):
|
||
t0 = time.time()
|
||
try:
|
||
resp = self._client.post(url, json=body, headers=headers)
|
||
elapsed = time.time() - t0
|
||
except httpx.RequestError as e:
|
||
last_error = f"network: {e}"
|
||
elapsed = time.time() - t0
|
||
self._log({
|
||
"tag": tag, "attempt": attempt, "elapsed": elapsed,
|
||
"error": last_error,
|
||
})
|
||
time.sleep(2 ** attempt)
|
||
continue
|
||
|
||
if resp.status_code == 200:
|
||
try:
|
||
data = resp.json()
|
||
except Exception as e:
|
||
raise ZenMuxError(f"invalid JSON from zenmux: {e}; body={resp.text[:500]}")
|
||
usage = data.get("usage", {}) or {}
|
||
with self._usage_lock:
|
||
self.usage.add(model, usage)
|
||
content = ""
|
||
choices = data.get("choices") or []
|
||
if choices:
|
||
msg = choices[0].get("message") or {}
|
||
content = msg.get("content") or ""
|
||
self._log({
|
||
"tag": tag, "model": model, "attempt": attempt,
|
||
"elapsed": round(elapsed, 2),
|
||
"usage": usage,
|
||
"out_chars": len(content),
|
||
"status": 200,
|
||
})
|
||
if not content.strip():
|
||
# zenmux 偶尔返 200 但 content 空；视作可重试
|
||
last_error = "empty content"
|
||
time.sleep(2 ** attempt)
|
||
continue
|
||
return content
|
||
|
||
# 非 200
|
||
retryable = resp.status_code in RETRYABLE_STATUSES
|
||
last_error = f"HTTP {resp.status_code}: {resp.text[:500]}"
|
||
self._log({
|
||
"tag": tag, "attempt": attempt, "elapsed": round(elapsed, 2),
|
||
"status": resp.status_code, "error": last_error,
|
||
"retryable": retryable,
|
||
})
|
||
if not retryable:
|
||
self.usage.failed_calls += 1
|
||
raise ZenMuxError(last_error)
|
||
sleep_for = min(60, (2 ** attempt) + (attempt * 0.5))
|
||
time.sleep(sleep_for)
|
||
|
||
self.usage.failed_calls += 1
|
||
raise ZenMuxError(f"max retries exhausted. last error: {last_error}")
|
||
|
||
|
||
def load_secrets(env_path: Path | None = None) -> None:
|
||
"""从 secrets.env 把 key 塞到 os.environ，便于脚本直接运行。
|
||
|
||
格式宽松：`KEY=VALUE` 每行一条，`#` 开头是注释，忽略空行。
|
||
"""
|
||
if env_path is None:
|
||
# 默认在仓库根找 secrets.env
|
||
here = Path(__file__).resolve()
|
||
for parent in [here.parent, *here.parents]:
|
||
cand = parent / "secrets.env"
|
||
if cand.exists():
|
||
env_path = cand
|
||
break
|
||
if not env_path or not env_path.exists():
|
||
return
|
||
for line in env_path.read_text(encoding="utf-8").splitlines():
|
||
line = line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
k, _, v = line.partition("=")
|
||
k = k.strip()
|
||
v = v.strip().strip('"').strip("'")
|
||
if k and v and k not in os.environ:
|
||
os.environ[k] = v
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 冒烟测试：python -m scripts.lib.zenmux_client
|
||
load_secrets()
|
||
with ZenMuxClient() as c:
|
||
out = c.chat_complete(
|
||
model="anthropic/claude-haiku-4.5",
|
||
system="You reply in exactly one English word.",
|
||
user="Say hello.",
|
||
max_tokens=20,
|
||
tag="smoke",
|
||
)
|
||
print("reply:", out)
|
||
print(c.usage.summary(), file=sys.stderr)
|