Files
ai-safety-platform/app/api/v1/endpoints/runs.py
T

559 lines
23 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import json
import logging
from hashlib import sha256
from typing import Annotated, Literal
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.export_service import matches_result, render_markdown
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. 每 12 秒调用 `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
]
@router.get(
"/runs/{run_id}/export",
summary="导出测试运行详细报告",
description="将终态 Run 的逐条结果导出为 Markdown;默认导出全部,可按执行状态和评价结果筛选。",
responses={
200: {"content": {"text/markdown": {}}, "description": "Markdown 详细报告附件。"},
404: {"model": ErrorResponse},
409: {"model": ErrorResponse},
},
)
def export_run(
run_id: int = Path(description="要导出的测试运行 ID。", gt=0),
execution_status: Literal["completed", "error"] | None = Query(default=None),
verdict: Literal["pass", "fail", "needs_human_review", "judge_format_error", "none"]
| None = Query(default=None),
db: Session = Depends(get_database),
settings: Settings = Depends(get_settings),
) -> Response:
run = TestRunRepository(db).get(run_id)
if not run:
raise NotFoundError("测试运行不存在")
if run.status not in TERMINAL_STATUSES:
raise ConflictError("只能导出已结束的测试运行")
results = [
item.model_dump()
for item in get_run_results(run_id, db, settings)
if matches_result(item.model_dump(), execution_status, verdict)
]
content = render_markdown(_run_detail(run).model_dump(mode="json"), results)
return Response(
content=content,
media_type="text/markdown",
headers={"Content-Disposition": f'attachment; filename="run_{run_id}_results.md"'},
)