baozaotumao2025
522 lines
21 KiB
Python
522 lines
21 KiB
Python
import json
|
||
import logging
|
||
from hashlib import sha256
|
||
from typing import Annotated
|
||
|
||
from fastapi import APIRouter, BackgroundTasks, Depends, Header, Path, Query, Response, status
|
||
|
||
from sqlalchemy.exc import IntegrityError
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.api.deps import get_database
|
||
from app.core.config import Settings, get_settings
|
||
from app.core.exceptions import ConflictError, NotFoundError
|
||
from app.db.repositories.repositories import ResultRepository, TestRunRepository
|
||
from app.db.session import SessionLocal
|
||
from app.providers.model_provider import ProviderFactory
|
||
from app.schemas.api import (
|
||
ErrorResponse,
|
||
ResultItem,
|
||
RetryErrorsResponse,
|
||
RetryResultResponse,
|
||
RunDetailResponse,
|
||
RunResponse,
|
||
ResumeRunResponse,
|
||
StartRunRequest,
|
||
UpdateRunRequest,
|
||
)
|
||
from app.services.capability_service import CapabilityProbeService
|
||
from app.services.dataset_service import DatasetGateway, TestPlanService
|
||
from app.services.test_execution_service import TestExecutionService
|
||
|
||
router = APIRouter()
|
||
logger = logging.getLogger(__name__)
|
||
active_run_ids: set[int] = set()
|
||
TERMINAL_STATUSES = {"cancelled", "completed", "completed_with_errors", "failed"}
|
||
|
||
|
||
def _run_detail(run) -> RunDetailResponse:
|
||
summary = json.loads(run.summary_json or "{}")
|
||
terminal = run.status in TERMINAL_STATUSES
|
||
processed = run.completed_count + run.error_count
|
||
return RunDetailResponse(
|
||
run_id=run.id,
|
||
status=run.status,
|
||
terminal=terminal,
|
||
phase=summary.get("phase", run.status),
|
||
selected_count=run.selected_count,
|
||
processed_count=processed,
|
||
completed_count=run.completed_count,
|
||
error_count=run.error_count,
|
||
progress_percent=(
|
||
round(processed / run.selected_count * 100, 2)
|
||
if run.selected_count
|
||
else (100.0 if terminal and run.status not in {"failed", "cancelled"} else 0.0)
|
||
),
|
||
current_execution_id=summary.get("current_execution_id"),
|
||
started_at=run.created_at,
|
||
finished_at=run.updated_at if terminal else None,
|
||
error_message=summary.get("error") if run.status == "failed" else None,
|
||
poll_after_seconds=0 if terminal else 2,
|
||
summary=summary,
|
||
)
|
||
|
||
|
||
async def execute_run(run_id: int, profile: str, auto_judge: bool, settings: Settings) -> None:
|
||
if run_id in active_run_ids:
|
||
return
|
||
active_run_ids.add(run_id)
|
||
try:
|
||
await _execute_run(run_id, profile, auto_judge, settings)
|
||
finally:
|
||
active_run_ids.discard(run_id)
|
||
|
||
|
||
async def _execute_run(run_id: int, profile: str, auto_judge: bool, settings: Settings) -> None:
|
||
with SessionLocal() as db:
|
||
repo = TestRunRepository(db)
|
||
run = repo.get(run_id)
|
||
if not run or run.status == "cancelled":
|
||
return
|
||
try:
|
||
factory = ProviderFactory(settings, db)
|
||
target = factory.create("target")
|
||
try:
|
||
capabilities = json.loads(run.capabilities_json) if run.capabilities_json else None
|
||
except json.JSONDecodeError:
|
||
logger.warning("Invalid capability cache; probing again run_id=%s", run_id)
|
||
capabilities = None
|
||
if capabilities is None:
|
||
repo.update_status(
|
||
run, "probing", summary={"phase": "probing", "auto_judge": auto_judge}
|
||
)
|
||
capabilities = await CapabilityProbeService(target).probe()
|
||
repo.save_capabilities(run, capabilities)
|
||
db.refresh(run)
|
||
if run.status == "cancelled":
|
||
return
|
||
plan = TestPlanService(DatasetGateway(settings)).build(
|
||
capabilities=capabilities,
|
||
profile=profile,
|
||
)
|
||
cases = plan["selected_cases"]
|
||
run.selected_count = len(cases)
|
||
repo.update_status(
|
||
run,
|
||
"running",
|
||
summary={
|
||
"phase": "executing",
|
||
"completed": 0,
|
||
"errors": 0,
|
||
"selected": len(cases),
|
||
"auto_judge": auto_judge,
|
||
},
|
||
)
|
||
await TestExecutionService(
|
||
db=db,
|
||
target_provider=target,
|
||
judge_provider=factory.create("judge") if auto_judge else None,
|
||
).execute(run=run, cases=cases, auto_judge=auto_judge)
|
||
except Exception as exc:
|
||
db.rollback()
|
||
run = repo.get(run_id)
|
||
if run:
|
||
message = str(exc) or type(exc).__name__
|
||
repo.update_status(
|
||
run,
|
||
"failed",
|
||
summary={"phase": "failed", "error": message, "auto_judge": auto_judge},
|
||
)
|
||
logger.exception("Background run failed run_id=%s", run_id)
|
||
|
||
|
||
@router.post(
|
||
"/runs",
|
||
status_code=status.HTTP_202_ACCEPTED,
|
||
response_model=RunResponse,
|
||
summary="启动安全测试",
|
||
description=(
|
||
"对当前生效的 `target` 模型提供商配置启动一次后台安全测试。"
|
||
"接口创建运行记录后立即返回 HTTP 202，能力探测、样例执行和裁判在后台进行。\n\n"
|
||
"### 请求示例\n\n"
|
||
'```json\n{\n "profile": "smoke",\n "auto_judge": true\n}\n```\n\n'
|
||
"* `profile=smoke`：单轮、多轮、工具和图片各选取 1 条；四类都可执行才启动。\n"
|
||
"* `profile=all`：执行当前模型能力支持的全部样例。\n"
|
||
"* `auto_judge=true`：每条回复调用当前生效的 `judge` 配置，输出 `pass`、`fail` 或 "
|
||
"`needs_human_review`；工具题会同时参考 `tool_call_policy` 与实际工具调用。"
|
||
"设为 `false` 则只保存被测模型原始回复。\n\n"
|
||
"### 执行与返回\n\n"
|
||
"返回 `202 Accepted` 与 `{run_id, status, selected_count}`。创建时 `status=pending`、"
|
||
"`selected_count=0`；能力筛选完成后才写入实际样例数。\n\n"
|
||
"### 幂等创建\n\n"
|
||
"客户端重试 `POST /runs` 时应携带 `Idempotency-Key` 请求头（最多 128 字符）。"
|
||
"相同键和相同请求参数返回同一个运行且不会重复投递后台任务；同一个键用于不同参数返回 409。"
|
||
"未提供该请求头时，每次调用都会创建一次新测试。\n\n"
|
||
"### 状态规则\n\n"
|
||
"| status | 含义 |\n| --- | --- |\n"
|
||
"| `pending` | 运行记录已创建，后台任务待启动 |\n"
|
||
"| `probing` | 正在探测目标模型能力 |\n"
|
||
"| `running` | 正在逐条执行样例及可选裁判 |\n"
|
||
"| `cancelled` | 已请求取消，当前上游请求结束后停止后续样例 |\n"
|
||
"| `completed` | 所有样例调用成功 |\n"
|
||
"| `completed_with_errors` | 整体已结束，但至少一条样例调用失败 |\n"
|
||
"| `failed` | 能力探测、数据集或全局配置阶段失败 |\n\n"
|
||
"### 重试规则\n\n"
|
||
"模型请求超时由 `REQUEST_TIMEOUT_SECONDS` 控制；网络异常、HTTP 408/429/5xx "
|
||
"会按 `REQUEST_RETRY_COUNT` 额外重试。例如配置为 5 时，最多发起 6 次请求。"
|
||
"HTTP 429、408 和 5xx 按指数退避等待，单次最多由 `REQUEST_RETRY_MAX_DELAY_SECONDS` 限制；"
|
||
"HTTP 429 的 `Retry-After` 头会优先采用。其他 4xx 不重试。\n\n"
|
||
"### 调用顺序\n\n"
|
||
"1. 每 1～2 秒调用 `GET /api/v1/runs/{run_id}` 查询进度。\n"
|
||
"2. 进入最终状态后停止轮询。\n"
|
||
"3. 调用 `GET /api/v1/runs/{run_id}/results` 查看逐条结果，或"
|
||
" `GET /api/v1/reports/{run_id}` 查看汇总。"
|
||
),
|
||
responses={
|
||
202: {"description": "运行已创建，后台任务已提交。"},
|
||
409: {"model": ErrorResponse, "description": "幂等键已用于不同请求参数。"},
|
||
500: {"model": ErrorResponse, "description": "运行记录创建失败。"},
|
||
},
|
||
)
|
||
async def start_run(
|
||
payload: StartRunRequest,
|
||
background_tasks: BackgroundTasks,
|
||
db: Session = Depends(get_database),
|
||
settings: Settings = Depends(get_settings),
|
||
idempotency_key: Annotated[str | None, Header(alias="Idempotency-Key", max_length=128)] = None,
|
||
):
|
||
repo = TestRunRepository(db)
|
||
fingerprint = sha256(
|
||
json.dumps(
|
||
{"profile": payload.profile, "auto_judge": payload.auto_judge},
|
||
sort_keys=True,
|
||
separators=(",", ":"),
|
||
).encode()
|
||
).hexdigest()
|
||
if idempotency_key:
|
||
existing = repo.get_by_idempotency_key(idempotency_key)
|
||
if existing:
|
||
if existing.request_fingerprint != fingerprint:
|
||
raise ConflictError("幂等键不能复用于不同请求")
|
||
return RunResponse(
|
||
run_id=existing.id,
|
||
status=existing.status,
|
||
selected_count=existing.selected_count,
|
||
)
|
||
try:
|
||
run = repo.create(
|
||
run_type="safety_test",
|
||
profile=payload.profile,
|
||
selected_count=0,
|
||
idempotency_key=idempotency_key,
|
||
request_fingerprint=fingerprint,
|
||
)
|
||
except IntegrityError:
|
||
db.rollback()
|
||
existing = repo.get_by_idempotency_key(idempotency_key or "")
|
||
if existing and existing.request_fingerprint == fingerprint:
|
||
return RunResponse(
|
||
run_id=existing.id,
|
||
status=existing.status,
|
||
selected_count=existing.selected_count,
|
||
)
|
||
raise
|
||
background_tasks.add_task(
|
||
execute_run,
|
||
run.id,
|
||
payload.profile,
|
||
payload.auto_judge,
|
||
settings,
|
||
)
|
||
return RunResponse(
|
||
run_id=run.id,
|
||
status=run.status,
|
||
selected_count=run.selected_count,
|
||
)
|
||
|
||
|
||
@router.get(
|
||
"/runs",
|
||
response_model=list[RunDetailResponse],
|
||
summary="列出测试运行",
|
||
description="按创建时间倒序返回测试运行；`limit` 默认 50，最大 200。",
|
||
)
|
||
def list_runs(
|
||
limit: int = Query(default=50, ge=1, le=200),
|
||
db: Session = Depends(get_database),
|
||
):
|
||
return [_run_detail(run) for run in TestRunRepository(db).list(limit)]
|
||
|
||
|
||
@router.post(
|
||
"/runs/{run_id}/resume",
|
||
status_code=status.HTTP_202_ACCEPTED,
|
||
response_model=ResumeRunResponse,
|
||
summary="从已保存结果恢复测试运行",
|
||
description=(
|
||
"仅用于已中断且确认原后台任务不再运行的测试。已持久化结果的样例不会再次请求模型，"
|
||
"包括 `execution_status=error` 的样例；该接口只补跑尚未落库的样例。若同一运行仍在当前进程执行，"
|
||
"重复调用仅返回当前运行，不会再次投递后台任务。`completed_with_errors` 应调用 "
|
||
"`retry-errors`，`completed` 无需恢复。"
|
||
),
|
||
responses={404: {"model": ErrorResponse}, 409: {"model": ErrorResponse}},
|
||
)
|
||
async def resume_run(
|
||
background_tasks: BackgroundTasks,
|
||
run_id: int = Path(description="要恢复的测试运行 ID，必须为正整数。", gt=0),
|
||
db: Session = Depends(get_database),
|
||
settings: Settings = Depends(get_settings),
|
||
):
|
||
run = TestRunRepository(db).get(run_id)
|
||
if not run:
|
||
raise NotFoundError("测试运行不存在")
|
||
skipped_count = len(ResultRepository(db).execution_ids_by_run(run_id))
|
||
if run_id in active_run_ids:
|
||
return ResumeRunResponse(
|
||
run_id=run.id,
|
||
status=run.status,
|
||
selected_count=run.selected_count,
|
||
skipped_count=skipped_count,
|
||
)
|
||
if run.status == "completed_with_errors":
|
||
raise ConflictError("该运行已有错误结果，请使用 retry-errors")
|
||
if run.status == "completed":
|
||
raise ConflictError("已完成运行无需恢复")
|
||
summary = json.loads(run.summary_json or "{}")
|
||
auto_judge = summary.get("auto_judge", True)
|
||
TestRunRepository(db).update_status(
|
||
run,
|
||
"pending",
|
||
summary={
|
||
"phase": "pending",
|
||
"resuming": True,
|
||
"skipped": skipped_count,
|
||
"auto_judge": auto_judge,
|
||
},
|
||
)
|
||
background_tasks.add_task(execute_run, run.id, run.profile, auto_judge, settings)
|
||
return ResumeRunResponse(
|
||
run_id=run.id,
|
||
status=run.status,
|
||
selected_count=run.selected_count,
|
||
skipped_count=skipped_count,
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/runs/{run_id}/retry-errors",
|
||
status_code=status.HTTP_202_ACCEPTED,
|
||
response_model=RetryErrorsResponse,
|
||
summary="重试运行中的错误样例",
|
||
description=(
|
||
"仅用于 `completed_with_errors`。删除该运行已落库的 `execution_status=error` 结果，"
|
||
"保留成功结果，并只重新执行失败样例。运行中或没有错误结果时返回 409。"
|
||
),
|
||
responses={404: {"model": ErrorResponse}, 409: {"model": ErrorResponse}},
|
||
)
|
||
async def retry_run_errors(
|
||
background_tasks: BackgroundTasks,
|
||
run_id: int = Path(description="要重试错误样例的测试运行 ID。", gt=0),
|
||
db: Session = Depends(get_database),
|
||
settings: Settings = Depends(get_settings),
|
||
):
|
||
run_repo = TestRunRepository(db)
|
||
run = run_repo.get(run_id)
|
||
if not run:
|
||
raise NotFoundError("测试运行不存在")
|
||
if run_id in active_run_ids or run.status != "completed_with_errors":
|
||
raise ConflictError("只有已结束且包含错误的运行可以重试")
|
||
result_repo = ResultRepository(db)
|
||
retry_count = result_repo.delete_errors_by_run(run_id)
|
||
if not retry_count:
|
||
db.rollback()
|
||
raise ConflictError("该运行没有可重试的 error 结果")
|
||
remaining = result_repo.list_by_run(run_id)
|
||
run.completed_count = sum(row.execution_status == "completed" for row in remaining)
|
||
run.error_count = 0
|
||
summary = json.loads(run.summary_json or "{}")
|
||
auto_judge = summary.get("auto_judge", True)
|
||
run_repo.update_status(
|
||
run,
|
||
"pending",
|
||
summary={
|
||
"phase": "pending",
|
||
"retry_errors": True,
|
||
"retry_count": retry_count,
|
||
"auto_judge": auto_judge,
|
||
},
|
||
)
|
||
background_tasks.add_task(execute_run, run.id, run.profile, auto_judge, settings)
|
||
logger.info("Run errors requeued run_id=%s retry_count=%s", run_id, retry_count)
|
||
return RetryErrorsResponse(
|
||
run_id=run.id,
|
||
status=run.status,
|
||
selected_count=run.selected_count,
|
||
retry_count=retry_count,
|
||
)
|
||
|
||
|
||
@router.post(
|
||
"/runs/{run_id}/results/{execution_id}/retry",
|
||
status_code=status.HTTP_202_ACCEPTED,
|
||
response_model=RetryResultResponse,
|
||
summary="重试运行中的指定样例",
|
||
description=(
|
||
"仅用于 `completed` 或 `completed_with_errors` 运行。删除指定 `execution_id` "
|
||
"的已有结果，保留其他结果，并使用原运行的 profile 和自动仲裁设置重新排队。"
|
||
"运行中、非终态或重复提交返回 409；运行或指定结果不存在返回 404。"
|
||
),
|
||
responses={404: {"model": ErrorResponse}, 409: {"model": ErrorResponse}},
|
||
)
|
||
async def retry_run_result(
|
||
background_tasks: BackgroundTasks,
|
||
run_id: int = Path(description="测试运行 ID。", gt=0),
|
||
execution_id: str = Path(description="要重试的样例执行 ID。", min_length=1, max_length=128),
|
||
db: Session = Depends(get_database),
|
||
settings: Settings = Depends(get_settings),
|
||
):
|
||
run_repo = TestRunRepository(db)
|
||
run = run_repo.get(run_id)
|
||
if not run:
|
||
raise NotFoundError("测试运行不存在")
|
||
if run_id in active_run_ids or run.status not in {"completed", "completed_with_errors"}:
|
||
raise ConflictError("只有已结束的运行可以重试指定样例")
|
||
result_repo = ResultRepository(db)
|
||
if not result_repo.delete_by_run_and_execution(run_id, execution_id):
|
||
db.rollback()
|
||
raise NotFoundError("测试结果不存在")
|
||
remaining = result_repo.list_by_run(run_id)
|
||
run.completed_count = sum(row.execution_status == "completed" for row in remaining)
|
||
run.error_count = sum(row.execution_status == "error" for row in remaining)
|
||
summary = json.loads(run.summary_json or "{}")
|
||
auto_judge = summary.get("auto_judge", True)
|
||
run_repo.update_status(
|
||
run,
|
||
"pending",
|
||
summary={
|
||
"phase": "pending",
|
||
"retry_execution_id": execution_id,
|
||
"auto_judge": auto_judge,
|
||
},
|
||
)
|
||
background_tasks.add_task(execute_run, run.id, run.profile, auto_judge, settings)
|
||
logger.info("Run result requeued run_id=%s execution_id=%s", run_id, execution_id)
|
||
return RetryResultResponse(
|
||
run_id=run.id,
|
||
status=run.status,
|
||
selected_count=run.selected_count,
|
||
execution_id=execution_id,
|
||
)
|
||
|
||
|
||
@router.get(
|
||
"/runs/{run_id}",
|
||
response_model=RunDetailResponse,
|
||
summary="查询测试运行状态",
|
||
description=(
|
||
"根据运行 ID 查询阶段和实时进度，建议每 1～2 秒轮询。\n\n"
|
||
"`selected_count` 是筛选后计划执行数；`completed_count` 只计成功产生结果的样例；"
|
||
"`error_count` 计调用目标模型或裁判模型失败的样例。执行阶段始终满足 "
|
||
"`completed_count + error_count <= selected_count`。\n\n"
|
||
"前端应以 `terminal` 为停止轮询的唯一判定；为 true 时后台已结束，"
|
||
"`poll_after_seconds` 为 0。`processed_count` 和 `progress_percent` 可直接用于进度条，"
|
||
"`current_execution_id` 用于展示当前样例。\n\n"
|
||
"`summary.phase` 可为 `probing`、`executing`、`finished` 或 `failed`。"
|
||
"当 phase 为 `failed` 时，`summary.error` 给出全局失败原因。运行 ID 不存在时返回 404。"
|
||
),
|
||
responses={404: {"model": ErrorResponse, "description": "指定的测试运行不存在。"}},
|
||
)
|
||
def get_run(
|
||
run_id: int = Path(description="要查询的测试运行 ID，必须为正整数。", gt=0),
|
||
db: Session = Depends(get_database),
|
||
):
|
||
run = TestRunRepository(db).get(run_id)
|
||
if not run:
|
||
raise NotFoundError("测试运行不存在")
|
||
return _run_detail(run)
|
||
|
||
|
||
@router.patch(
|
||
"/runs/{run_id}",
|
||
response_model=RunDetailResponse,
|
||
summary="取消测试运行",
|
||
responses={404: {"model": ErrorResponse}, 409: {"model": ErrorResponse}},
|
||
)
|
||
def update_run(
|
||
payload: UpdateRunRequest,
|
||
run_id: int = Path(description="测试运行 ID。", gt=0),
|
||
db: Session = Depends(get_database),
|
||
):
|
||
repo = TestRunRepository(db)
|
||
run = repo.get(run_id)
|
||
if not run:
|
||
raise NotFoundError("测试运行不存在")
|
||
if run.status in TERMINAL_STATUSES:
|
||
raise ConflictError("测试运行已处于终态")
|
||
repo.update_status(run, payload.status, summary={"phase": "cancelled", "cancelled": True})
|
||
logger.info("Run cancelled run_id=%s", run_id)
|
||
return _run_detail(run)
|
||
|
||
|
||
@router.delete(
|
||
"/runs/{run_id}",
|
||
status_code=status.HTTP_204_NO_CONTENT,
|
||
summary="删除测试运行",
|
||
responses={404: {"model": ErrorResponse}, 409: {"model": ErrorResponse}},
|
||
)
|
||
def delete_run(
|
||
run_id: int = Path(description="测试运行 ID。", gt=0),
|
||
db: Session = Depends(get_database),
|
||
) -> Response:
|
||
repo = TestRunRepository(db)
|
||
run = repo.get(run_id)
|
||
if not run:
|
||
raise NotFoundError("测试运行不存在")
|
||
if run.status not in TERMINAL_STATUSES:
|
||
raise ConflictError("只能删除已结束或已取消的测试运行")
|
||
repo.delete(run)
|
||
logger.info("Run deleted run_id=%s", run_id)
|
||
return Response(status_code=status.HTTP_204_NO_CONTENT)
|
||
|
||
|
||
@router.get(
|
||
"/runs/{run_id}/results",
|
||
response_model=list[ResultItem],
|
||
summary="查询测试结果",
|
||
description=(
|
||
"返回已持久化的逐样例结果；运行中调用时可能只返回部分列表。\n\n"
|
||
"`execution_status=error` 表示该样例调用失败，原因在 `error_message`；"
|
||
"`verdict` 是独立的仲裁结论。"
|
||
"`model_input` 是与 `execution_id` 对应的完整测试输入，包含 messages、system 上下文及可选工具定义。"
|
||
"开启自动裁判时，成功样例的 status 通常为 `pass`、`fail` 或 "
|
||
"`needs_human_review`，结构化裁判放在 `judge_result`；关闭时 status 为 `completed`、"
|
||
"`judge_result` 为空对象。未知运行 ID 返回 404。"
|
||
),
|
||
)
|
||
def get_run_results(
|
||
run_id: int = Path(description="要查询结果的测试运行 ID，必须为正整数。", gt=0),
|
||
db: Session = Depends(get_database),
|
||
settings: Settings = Depends(get_settings),
|
||
):
|
||
if not TestRunRepository(db).get(run_id):
|
||
raise NotFoundError("测试运行不存在")
|
||
rows = ResultRepository(db).list_by_run(run_id)
|
||
inputs = {
|
||
case["execution_id"]: case.get("model_input", {})
|
||
for case in DatasetGateway(settings).load()
|
||
}
|
||
return [
|
||
ResultItem(
|
||
execution_id=row.execution_id,
|
||
case_kind=row.case_kind,
|
||
interaction_mode=row.interaction_mode,
|
||
execution_status=row.execution_status,
|
||
verdict=row.verdict,
|
||
model_input=inputs.get(row.execution_id, {}),
|
||
model_response=row.model_response,
|
||
judge_result=json.loads(row.judge_result_json or "{}"),
|
||
error_message=row.error_message,
|
||
)
|
||
for row in rows
|
||
]
|