ai-safety-platform/scripts/export_run.py

#!/usr/bin/env python3
"""Export one completed safety-test run as a readable Markdown report."""

import argparse
import getpass
import json
import re
import sys
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


def request_json(
    base_url: str,
    path: str,
    *,
    token: str | None = None,
    payload: dict[str, Any] | None = None,
) -> Any:
    headers = {"Accept": "application/json"}
    data = None
    if payload is not None:
        headers["Content-Type"] = "application/json"
        data = json.dumps(payload).encode()
    if token:
        headers["Authorization"] = f"Bearer {token}"
    request = Request(base_url.rstrip("/") + path, data=data, headers=headers)
    try:
        with urlopen(request, timeout=30) as response:
            return json.load(response)
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        try:
            message = json.loads(body).get("error", {}).get("message", body)
        except json.JSONDecodeError:
            message = body
        raise RuntimeError(f"HTTP {exc.code}: {message}") from exc
    except URLError as exc:
        raise RuntimeError(f"无法连接服务：{exc.reason}") from exc


def block(value: Any) -> str:
    text = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, indent=2)
    longest = max((len(part) for part in re.findall(r"`+", text)), default=0)
    fence = "`" * max(3, longest + 1)
    return f"{fence}\n{text}\n{fence}"


def render_input(model_input: dict[str, Any]) -> list[str]:
    sections: list[str] = []
    if system := model_input.get("system"):
        sections.extend(["#### System", "", block(system), ""])
    for message in model_input.get("messages", []):
        role = str(message.get("role", "unknown")).capitalize()
        sections.extend([f"#### {role}", "", block(message.get("content", "")), ""])
    extras = {
        key: value
        for key, value in model_input.items()
        if key not in {"system", "messages"} and value is not None and value != "" and value is not False
    }
    if extras:
        sections.extend(["#### 其他输入信息", "", block(extras), ""])
    return sections or ["无", ""]


def is_completed_pass(result: dict[str, Any]) -> bool:
    return result.get("execution_status") == "completed" and result.get("verdict") == "pass"


def render_markdown(run: dict[str, Any], results: list[dict[str, Any]]) -> str:
    lines = [
        f"# Run {run['run_id']} 安全测试详细报告",
        "",
        "## 运行概况",
        "",
        f"- 状态：{run.get('status', 'unknown')}",
        f"- 阶段：{run.get('phase', 'unknown')}",
        f"- 样例总数：{run.get('selected_count', len(results))}",
        f"- 已完成：{run.get('completed_count', 0)}",
        f"- 执行错误：{run.get('error_count', 0)}",
        f"- 开始时间：{run.get('started_at', '')}",
        f"- 完成时间：{run.get('finished_at', '')}",
        "",
        f"- 本报告 Pass 样例：{len(results)}",
        "",
        "## Pass 样例详情",
        "",
    ]
    for index, result in enumerate(results, 1):
        judge = result.get("judge_result") or {}
        lines.extend(
            [
                f"### {index}. {result.get('execution_id', 'unknown')}",
                "",
                f"- 测试类型：{result.get('case_kind', 'unknown')}",
                f"- 交互模式：{result.get('interaction_mode', 'unknown')}",
                f"- 执行状态：{result.get('execution_status', 'unknown')}",
                f"- 仲裁结论：{result.get('verdict') or '未仲裁'}",
                "",
                "#### 模型输入",
                "",
                *render_input(result.get("model_input") or {}),
                "#### 模型输出",
                "",
                block(result.get("model_response", "")),
                "",
                "#### 仲裁结果",
                "",
            ]
        )
        if judge:
            lines.extend(
                [
                    f"- 结论：{judge.get('verdict', result.get('verdict', 'unknown'))}",
                    f"- 评分：{judge.get('score', '无')}",
                    f"- 理由：{judge.get('reason', '无')}",
                ]
            )
            if raw := judge.get("raw_response"):
                lines.extend(["", "##### 仲裁原始输出", "", block(raw)])
        else:
            lines.append("未启用自动仲裁或没有仲裁结果。")
        if error := result.get("error_message"):
            lines.extend(["", "#### 执行错误", "", block(error)])
        lines.extend(["", "---", ""])
    return "\n".join(lines).rstrip() + "\n"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="将已完成的 run 导出为 Markdown 报告")
    parser.add_argument("--run-id", type=int, default=7, help="要导出的 run ID（默认：7）")
    parser.add_argument("--base-url", default="http://127.0.0.1:8000", help="API 服务地址")
    parser.add_argument("--username", help="登录用户名；省略时交互输入")
    parser.add_argument("--output", type=Path, help="输出文件路径")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.run_id < 1:
        print("错误：--run-id 必须是正整数", file=sys.stderr)
        return 2
    username = args.username or input("Username: ").strip()
    password = getpass.getpass("Password: ")
    try:
        login = request_json(
            args.base_url,
            "/api/v1/auth/login",
            payload={"username": username, "password": password},
        )
        token = login["access_token"]
        run = request_json(args.base_url, f"/api/v1/runs/{args.run_id}", token=token)
        if not run.get("terminal"):
            raise RuntimeError(f"Run {args.run_id} 尚未结束，当前状态：{run.get('status')}")
        results = request_json(
            args.base_url, f"/api/v1/runs/{args.run_id}/results", token=token
        )
        results = [result for result in results if is_completed_pass(result)]
        output = args.output or Path("outputs") / f"run_{args.run_id}_pass_results.md"
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(render_markdown(run, results), encoding="utf-8")
    except (RuntimeError, KeyError, OSError) as exc:
        print(f"导出失败：{exc}", file=sys.stderr)
        return 1
    print(f"导出完成：{output}（共 {len(results)} 个测试样例）")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())