223 lines
7.9 KiB
Python
223 lines
7.9 KiB
Python
"""zcbot `/v1` 黑盒客户端。"""
|
||
from __future__ import annotations
|
||
|
||
import ipaddress
|
||
import os
|
||
import time
|
||
from typing import Any
|
||
from urllib.parse import urljoin, urlparse
|
||
|
||
import httpx
|
||
|
||
from .models import RunObservation
|
||
|
||
|
||
class EvalClientError(RuntimeError):
|
||
"""zcbot API 调用失败。"""
|
||
|
||
|
||
def require_safe_base_url(base_url: str, *, allow_remote: bool) -> str:
|
||
normalized = base_url.rstrip("/") + "/"
|
||
parsed = urlparse(normalized)
|
||
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
||
raise EvalClientError(f"无效 base_url: {base_url!r}")
|
||
if allow_remote:
|
||
return normalized
|
||
host = parsed.hostname.lower()
|
||
safe = host in {"localhost", "127.0.0.1", "::1"}
|
||
if not safe:
|
||
try:
|
||
safe = ipaddress.ip_address(host).is_loopback
|
||
except ValueError:
|
||
safe = False
|
||
if not safe:
|
||
raise EvalClientError(
|
||
f"默认只允许本机 zcbot，当前是 {host!r}；确认测试环境后传 --allow-remote"
|
||
)
|
||
return normalized
|
||
|
||
|
||
def _response_text(payload: Any) -> str:
|
||
if isinstance(payload, str):
|
||
return payload
|
||
if isinstance(payload, list):
|
||
parts: list[str] = []
|
||
for item in payload:
|
||
if isinstance(item, dict) and isinstance(item.get("text"), str):
|
||
parts.append(item["text"])
|
||
elif isinstance(item, str):
|
||
parts.append(item)
|
||
return "\n".join(parts)
|
||
return "" if payload is None else str(payload)
|
||
|
||
|
||
class ZcbotClient:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
base_url: str,
|
||
token: str,
|
||
allow_remote: bool = False,
|
||
request_timeout_s: float = 30.0,
|
||
) -> None:
|
||
self.base_url = require_safe_base_url(base_url, allow_remote=allow_remote)
|
||
if not token.strip():
|
||
raise EvalClientError("缺少 token；请设置配置中 token_env 指向的环境变量")
|
||
self._client = httpx.Client(
|
||
base_url=self.base_url,
|
||
headers={"Authorization": f"Bearer {token.strip()}"},
|
||
timeout=request_timeout_s,
|
||
follow_redirects=False,
|
||
# 评测 JWT 不得被 HTTP(S)_PROXY 透明转发；部署内网应直连测试实例。
|
||
trust_env=False,
|
||
)
|
||
|
||
@classmethod
|
||
def from_config(
|
||
cls, raw: dict[str, Any], *, allow_remote: bool = False
|
||
) -> "ZcbotClient":
|
||
token_env = str(raw.get("token_env", "ZCBOT_EVAL_TOKEN"))
|
||
return cls(
|
||
base_url=str(raw.get("base_url", "http://127.0.0.1:8765")),
|
||
token=os.getenv(token_env, ""),
|
||
allow_remote=allow_remote,
|
||
request_timeout_s=float(raw.get("request_timeout_s", 30)),
|
||
)
|
||
|
||
def close(self) -> None:
|
||
self._client.close()
|
||
|
||
def __enter__(self) -> "ZcbotClient":
|
||
return self
|
||
|
||
def __exit__(self, *_args: object) -> None:
|
||
self.close()
|
||
|
||
def _json(self, method: str, path: str, **kwargs: Any) -> dict[str, Any]:
|
||
response = self._client.request(method, path, **kwargs)
|
||
if response.status_code >= 400:
|
||
raise EvalClientError(
|
||
f"{method} {path} -> {response.status_code}: {response.text[:500]}"
|
||
)
|
||
try:
|
||
body = response.json()
|
||
except ValueError as exc:
|
||
raise EvalClientError(f"{method} {path} 返回了非 JSON") from exc
|
||
if not isinstance(body, dict):
|
||
raise EvalClientError(f"{method} {path} 返回 JSON 顶层不是 object")
|
||
return body
|
||
|
||
def create_task(
|
||
self,
|
||
*,
|
||
name: str,
|
||
working_dir: str,
|
||
skill: str,
|
||
model_profile: str,
|
||
) -> dict[str, Any]:
|
||
return self._json(
|
||
"POST",
|
||
"/v1/tasks",
|
||
json={
|
||
"name": name,
|
||
"working_dir": working_dir,
|
||
"description": "zcbot evaluation task",
|
||
"skill": skill,
|
||
"model_profile": model_profile,
|
||
},
|
||
)
|
||
|
||
def run_prompt(
|
||
self,
|
||
*,
|
||
task_id: str,
|
||
prompt: str,
|
||
working_dir: str,
|
||
timeout_s: float,
|
||
) -> RunObservation:
|
||
started = time.monotonic()
|
||
body = self._json(
|
||
"POST", f"/v1/tasks/{task_id}/messages", json={"content": prompt}
|
||
)
|
||
events_url = str(body.get("events_url", ""))
|
||
if not events_url:
|
||
raise EvalClientError("POST messages 未返回 events_url")
|
||
stream_error = self._consume_events(events_url, timeout_s=timeout_s)
|
||
# AgentLoop 的 done 事件可能早于 BG worker 最终把 task.run_status 写回 idle；
|
||
# 短暂轮询终态，避免把已成功任务误记为 running。
|
||
terminal_deadline = time.monotonic() + min(15.0, timeout_s)
|
||
while True:
|
||
task = self._json("GET", f"/v1/tasks/{task_id}")
|
||
if task.get("run_status") not in {"running", "cancelling"}:
|
||
break
|
||
if time.monotonic() >= terminal_deadline:
|
||
break
|
||
time.sleep(0.1)
|
||
messages = self._json("GET", f"/v1/tasks/{task_id}/messages")
|
||
response = ""
|
||
for item in reversed(messages.get("messages", [])):
|
||
payload = item.get("payload", {}) if isinstance(item, dict) else {}
|
||
if isinstance(payload, dict) and payload.get("role") == "assistant":
|
||
response = _response_text(payload.get("content"))
|
||
break
|
||
run_error = str(task.get("run_error") or stream_error or "")
|
||
return RunObservation(
|
||
response=response,
|
||
duration_s=time.monotonic() - started,
|
||
cost_cny=float(task.get("cost_cny") or 0),
|
||
run_status=str(task.get("run_status") or ""),
|
||
run_error=run_error,
|
||
task_id=task_id,
|
||
working_dir=working_dir,
|
||
model_profile=str(task.get("model_profile") or ""),
|
||
artifact_loader=lambda relative: self.download_artifact(
|
||
working_dir, relative
|
||
),
|
||
)
|
||
|
||
def _consume_events(self, path: str, *, timeout_s: float) -> str:
|
||
deadline = time.monotonic() + timeout_s
|
||
error = ""
|
||
timeout = httpx.Timeout(
|
||
connect=min(30.0, timeout_s),
|
||
read=max(30.0, timeout_s),
|
||
write=30.0,
|
||
pool=30.0,
|
||
)
|
||
with self._client.stream("GET", path, timeout=timeout) as response:
|
||
if response.status_code >= 400:
|
||
raise EvalClientError(
|
||
f"GET {path} -> {response.status_code}: {response.read()[:500]!r}"
|
||
)
|
||
event = ""
|
||
data = ""
|
||
for line in response.iter_lines():
|
||
if time.monotonic() > deadline:
|
||
raise EvalClientError(f"任务运行超过 {timeout_s:.0f}s")
|
||
if line.startswith("event:"):
|
||
event = line[6:].strip()
|
||
elif line.startswith("data:"):
|
||
data = line[5:].strip()
|
||
elif not line:
|
||
if event == "error":
|
||
error = data[:1000]
|
||
if event in {"done", "error"}:
|
||
return error
|
||
event, data = "", ""
|
||
return error
|
||
|
||
def download_artifact(self, working_dir: str, relative_path: str) -> bytes:
|
||
relative = relative_path.replace("\\", "/").strip("/")
|
||
if not relative or ".." in relative.split("/"):
|
||
raise FileNotFoundError(relative_path)
|
||
path = f"{working_dir.strip('/')}/{relative}"
|
||
response = self._client.get("/v1/files/download", params={"path": path})
|
||
if response.status_code == 404:
|
||
raise FileNotFoundError(relative_path)
|
||
if response.status_code >= 400:
|
||
raise EvalClientError(
|
||
f"GET /v1/files/download {path!r} -> {response.status_code}: "
|
||
f"{response.text[:500]}"
|
||
)
|
||
return response.content
|