fix(backend): 落实 PR #9 评审意见
- 检索调优参数(rrf_k/rerank/rerank_candidates/score_threshold)透传到引擎实际执行 - Recall 去重，避免同一 Note 多 Block 重复导致 Recall 超 1 - RAG 运行改为后台异步执行：创建即 queued + 202，支持取消与 SSE 实时事件 - 数据集元数据校验，坏文件隔离跳过；citation_required 语义修正 - modes 空/重复校验；配置快照记录模型版本与索引元信息 Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
Claude Code
@@ -1,19 +1,22 @@
|
||||
"""Benchmark 服务：运行注册表、配置快照与报告组装。
|
||||
|
||||
MVP 阶段运行是同步的（与 index_service 一致）：POST 创建后立即执行完并返回
|
||||
completed 的 BenchmarkRun。运行记录、事件与报告暂存内存（_runs/_events/_reports），
|
||||
不持久化到 SQLite；后续接入异步任务队列时再落库。
|
||||
RAG Benchmark 采用「创建即返回 queued、后台 Task 异步执行」的模式（与 index_service
|
||||
的 rebuild 一致）：POST 创建后立即返回 202 queued 的 BenchmarkRun，由受管 asyncio.Task
|
||||
在后台逐 Case 求值，进度与事件实时写入内存注册表，供 SSE 订阅。运行记录、事件与报告
|
||||
暂存内存（_runs/_events/_reports），不持久化到 SQLite；后续接入异步任务队列时再落库。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from uuid import uuid4
|
||||
|
||||
from app import repository
|
||||
from app.benchmarks import datasets
|
||||
from app.benchmarks.datasets import RAGDataset
|
||||
from app.benchmarks.rag import run_rag
|
||||
from app.benchmarks.rag import BenchmarkCancelled, run_rag
|
||||
from app.config import get_settings
|
||||
from app.contracts import (
|
||||
BenchmarkEvent,
|
||||
@@ -32,6 +35,9 @@ from app.retrieval.engine import engine
|
||||
_runs: dict[str, BenchmarkRun] = {}
|
||||
_events: dict[str, list[BenchmarkEvent]] = {}
|
||||
_reports: dict[str, BenchmarkReport] = {}
|
||||
_tasks: dict[str, asyncio.Task] = {}
|
||||
_subscribers: dict[str, list[asyncio.Queue[BenchmarkEvent]]] = {}
|
||||
_cancel_flags: dict[str, asyncio.Event] = {}
|
||||
MAX_RUNS = 100
|
||||
|
||||
|
||||
@@ -46,6 +52,9 @@ def _remember(run: BenchmarkRun) -> None:
|
||||
_runs.pop(oldest, None)
|
||||
_events.pop(oldest, None)
|
||||
_reports.pop(oldest, None)
|
||||
_tasks.pop(oldest, None)
|
||||
_subscribers.pop(oldest, None)
|
||||
_cancel_flags.pop(oldest, None)
|
||||
|
||||
|
||||
def _config_snapshot(request: RAGRunRequest, dataset: RAGDataset) -> dict:
|
||||
@@ -58,8 +67,16 @@ def _config_snapshot(request: RAGRunRequest, dataset: RAGDataset) -> dict:
|
||||
"modes": [m.value for m in request.modes],
|
||||
"retrieval": request.retrieval.model_dump(),
|
||||
"repeat": request.repeat,
|
||||
"embedding": {"model_id": engine.embedding.model_id, "dim": engine.embedding.dim},
|
||||
"reranker": {"model_id": engine.reranker.model_id},
|
||||
"embedding": {
|
||||
"model_id": engine.embedding.model_id,
|
||||
"version": engine.embedding.version,
|
||||
"dim": engine.embedding.dim,
|
||||
},
|
||||
"reranker": {
|
||||
"model_id": engine.reranker.model_id,
|
||||
"version": engine.reranker.version,
|
||||
},
|
||||
"index_meta": repository.get_index_meta(),
|
||||
"app": {"version": settings.version, "environment": settings.environment},
|
||||
"python": sys.version.split()[0],
|
||||
"metadata": request.metadata,
|
||||
@@ -67,7 +84,7 @@ def _config_snapshot(request: RAGRunRequest, dataset: RAGDataset) -> dict:
|
||||
|
||||
|
||||
async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
|
||||
"""创建并同步执行一次 RAG Benchmark，返回 completed 的 BenchmarkRun。"""
|
||||
"""创建一次 RAG Benchmark，立即返回 queued 的 BenchmarkRun，由后台 Task 执行。"""
|
||||
dataset = datasets.load_dataset(request.dataset_id, BenchmarkKind.rag)
|
||||
run_id = "benchmark_" + uuid4().hex[:12]
|
||||
snapshot = _config_snapshot(request, dataset)
|
||||
@@ -77,24 +94,41 @@ async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
|
||||
kind=BenchmarkKind.rag,
|
||||
dataset_id=dataset.dataset_id,
|
||||
dataset_hash=dataset.content_hash,
|
||||
status=BenchmarkStatus.running,
|
||||
status=BenchmarkStatus.queued,
|
||||
progress=0.0,
|
||||
config_snapshot=snapshot,
|
||||
created_at=_now(),
|
||||
started_at=_now(),
|
||||
)
|
||||
_remember(run)
|
||||
_events[run_id] = []
|
||||
_subscribers[run_id] = []
|
||||
_cancel_flags[run_id] = asyncio.Event()
|
||||
_tasks[run_id] = asyncio.create_task(_execute_rag(run_id, request, dataset, snapshot))
|
||||
return run
|
||||
|
||||
|
||||
async def _execute_rag(
|
||||
run_id: str, request: RAGRunRequest, dataset: RAGDataset, snapshot: dict
|
||||
) -> None:
|
||||
"""后台执行 RAG Benchmark，实时更新进度/事件，结束后写入报告并关闭订阅。"""
|
||||
cancel_event = _cancel_flags[run_id]
|
||||
|
||||
def emit(event_type: BenchmarkEventType, data: dict) -> None:
|
||||
sequence = len(_events[run_id])
|
||||
_events[run_id].append(
|
||||
BenchmarkEvent(
|
||||
event=event_type, run_id=run_id, sequence=sequence,
|
||||
data=data, timestamp=_now(),
|
||||
)
|
||||
event = BenchmarkEvent(
|
||||
event=event_type, run_id=run_id, sequence=sequence, data=data, timestamp=_now()
|
||||
)
|
||||
_events[run_id].append(event)
|
||||
for queue in _subscribers.get(run_id, []):
|
||||
queue.put_nowait(event)
|
||||
|
||||
def finish() -> None:
|
||||
_subscribers.pop(run_id, None)
|
||||
_cancel_flags.pop(run_id, None)
|
||||
|
||||
_runs[run_id] = _runs[run_id].model_copy(
|
||||
update={"status": BenchmarkStatus.running, "started_at": _now()}
|
||||
)
|
||||
emit(
|
||||
BenchmarkEventType.run_started,
|
||||
{"dataset_id": dataset.dataset_id, "modes": [m.value for m in request.modes]},
|
||||
@@ -107,8 +141,31 @@ async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
|
||||
emit(BenchmarkEventType.case_completed, result.model_dump(mode="json"))
|
||||
|
||||
try:
|
||||
metrics_by_mode, results = await run_rag(dataset, request, on_case=on_case)
|
||||
except Exception as exc:
|
||||
metrics_by_mode, results = await run_rag(
|
||||
dataset,
|
||||
request,
|
||||
on_case=on_case,
|
||||
should_cancel=cancel_event.is_set,
|
||||
)
|
||||
except BenchmarkCancelled:
|
||||
_runs[run_id] = _runs[run_id].model_copy(
|
||||
update={
|
||||
"status": BenchmarkStatus.cancelled,
|
||||
"progress": 1.0,
|
||||
"completed_at": _now(),
|
||||
}
|
||||
)
|
||||
_reports[run_id] = BenchmarkReport(
|
||||
run_id=run_id,
|
||||
kind=BenchmarkKind.rag,
|
||||
dataset_id=dataset.dataset_id,
|
||||
dataset_hash=dataset.content_hash,
|
||||
status=BenchmarkStatus.cancelled,
|
||||
config_snapshot=snapshot,
|
||||
)
|
||||
finish()
|
||||
return
|
||||
except Exception as exc: # 单次运行失败不拖垮服务，记录错误后结束
|
||||
_runs[run_id] = _runs[run_id].model_copy(
|
||||
update={
|
||||
"status": BenchmarkStatus.failed,
|
||||
@@ -127,7 +184,8 @@ async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
|
||||
config_snapshot=snapshot,
|
||||
error=str(exc),
|
||||
)
|
||||
raise ApiError(500, "BENCHMARK_RUN_FAILED", str(exc), {"run_id": run_id}) from exc
|
||||
finish()
|
||||
return
|
||||
|
||||
metrics = {mode: m.model_dump() for mode, m in metrics_by_mode.items()}
|
||||
_runs[run_id] = _runs[run_id].model_copy(
|
||||
@@ -149,7 +207,7 @@ async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
|
||||
metrics=metrics,
|
||||
cases=results,
|
||||
)
|
||||
return _runs[run_id]
|
||||
finish()
|
||||
|
||||
|
||||
def list_runs(
|
||||
@@ -181,13 +239,38 @@ def get_events(run_id: str) -> list[BenchmarkEvent]:
|
||||
|
||||
|
||||
def cancel_run(run_id: str) -> BenchmarkRun | None:
|
||||
"""取消运行：同步 MVP 下运行通常已结束，仅对仍在排队/运行的记录置为 cancelled。"""
|
||||
"""取消运行：对 queued/running 设置取消标志，后台 Task 在 Case 边界检查后置为 cancelled。"""
|
||||
run = _runs.get(run_id)
|
||||
if run is None:
|
||||
return None
|
||||
if run.status in (BenchmarkStatus.queued, BenchmarkStatus.running):
|
||||
run = run.model_copy(
|
||||
update={"status": BenchmarkStatus.cancelled, "completed_at": _now()}
|
||||
)
|
||||
_runs[run_id] = run
|
||||
_cancel_flags[run_id].set()
|
||||
return run
|
||||
|
||||
|
||||
def subscribe(run_id: str) -> asyncio.Queue[BenchmarkEvent] | None:
|
||||
"""订阅运行事件流；运行已结束（completed/failed/cancelled）时返回 None。"""
|
||||
run = _runs.get(run_id)
|
||||
if run is None or run.status in (
|
||||
BenchmarkStatus.completed,
|
||||
BenchmarkStatus.failed,
|
||||
BenchmarkStatus.cancelled,
|
||||
):
|
||||
return None
|
||||
queue: asyncio.Queue[BenchmarkEvent] = asyncio.Queue()
|
||||
_subscribers.setdefault(run_id, []).append(queue)
|
||||
return queue
|
||||
|
||||
|
||||
def unsubscribe(run_id: str, queue: asyncio.Queue[BenchmarkEvent]) -> None:
|
||||
subscribers = _subscribers.get(run_id)
|
||||
if subscribers and queue in subscribers:
|
||||
subscribers.remove(queue)
|
||||
|
||||
|
||||
async def wait_for_run(run_id: str) -> BenchmarkRun:
|
||||
"""等待后台任务结束（测试/轮询用）；无任务时直接返回当前状态。"""
|
||||
task = _tasks.get(run_id)
|
||||
if task is not None:
|
||||
await task
|
||||
return _runs.get(run_id)
|
||||
|
||||
Reference in New Issue
Block a user