v0.12: stabilize search routing and profile-driven phase4 pipeline

scripts/lib/zenmux_client.py

@@ -141,6 +141,8 @@ class ZenMuxClient:
         temperature: float = 0.3,
         max_tokens: int = 16000,
         extra_messages: list[dict[str, str]] | None = None,
+        web_search: bool = False,
+        web_search_options: dict[str, Any] | None = None,
         tag: str = "",
     ) -> str:
         """一次非流式对话补全。
@@ -167,6 +169,8 @@ class ZenMuxClient:
             "temperature": temperature,
             "max_tokens": max_tokens,
         }
+        if web_search:
+            body["web_search_options"] = web_search_options or {}
         headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
@@ -233,6 +237,116 @@ class ZenMuxClient:
         self.usage.failed_calls += 1
         raise ZenMuxError(f"max retries exhausted. last error: {last_error}")
 
+    def chat_complete_with_meta(
+        self,
+        model: str,
+        system: str,
+        user: str,
+        *,
+        temperature: float = 0.3,
+        max_tokens: int = 16000,
+        extra_messages: list[dict[str, str]] | None = None,
+        web_search: bool = False,
+        web_search_options: dict[str, Any] | None = None,
+        tag: str = "",
+    ) -> dict[str, Any]:
+        """Return content plus metadata from one completion call."""
+        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,
+        }
+        if web_search:
+            body["web_search_options"] = web_search_options or {}
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+        url = f"{self.base_url}/chat/completions"
+
+        last_error = ""
+        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:
+                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)
+                continue
+
+            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)
+
+            message = ((data.get("choices") or [{}])[0].get("message") or {})
+            content = message.get("content") or ""
+            annotations = message.get("annotations") or []
+            urls: list[str] = []
+            for ann in annotations:
+                if not isinstance(ann, dict):
+                    continue
+                citation = ann.get("url_citation") or {}
+                url_item = citation.get("url")
+                if url_item:
+                    urls.append(url_item)
+            self._log({
+                "tag": tag,
+                "model": model,
+                "attempt": attempt,
+                "elapsed": round(elapsed, 2),
+                "usage": usage,
+                "out_chars": len(content),
+                "status": 200,
+                "web_search": web_search,
+                "citations": len(urls),
+            })
+            if not content.strip():
+                last_error = "empty content"
+                time.sleep(2 ** attempt)
+                continue
+            return {
+                "content": content,
+                "usage": usage,
+                "citations": urls,
+                "raw": data,
+            }
+
+        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，便于脚本直接运行。