"""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 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) 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_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 {} 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)