feat: 完成benchmark后端功能

This commit is contained in:
yxx
2026-09-01 23:19:28 +08:00
parent 83782f1d0a
commit 9b50b8f0ce
11 changed files with 1060 additions and 1 deletions
+8
View File
@@ -0,0 +1,8 @@
"""Benchmark 服务:RAG / Agent 数据集注册、指标计算与运行管理。
模块划分:
- metrics.py 纯函数指标(Hit@K / Recall@K / MRR / CitationHit / 分位数)
- datasets.py 受控目录的 Dataset 注册与校验
- rag.py RAG Benchmark Runner(调用 retrieval.engine.search
- service.py 运行注册表、配置快照与报告组装
"""
+162
View File
@@ -0,0 +1,162 @@
"""Benchmark Dataset 注册:从受控目录加载 JSON 数据集并校验。
Dataset 只能来自配置目录(settings.benchmark_datasets_path),API 不接受调用方提交
任意文件路径。目录不存在或为空时按「无数据集」处理,不报错。
"""
from __future__ import annotations
import hashlib
import json
from dataclasses import dataclass, field
from pathlib import Path
from pydantic import ValidationError
from app.config import get_settings
from app.contracts import (
BenchmarkDatasetInfo,
BenchmarkKind,
RAGDatasetCase,
)
from app.errors import ApiError
@dataclass
class RAGDataset:
"""内存中的 RAG 数据集:元信息 + 已校验的 Case 列表 + 内容哈希。"""
dataset_id: str
kind: BenchmarkKind
version: str
description: str
cases: list[RAGDatasetCase] = field(default_factory=list)
content_hash: str = ""
def _datasets_dir() -> Path:
return get_settings().benchmark_datasets_path
def _dataset_files() -> list[Path]:
directory = _datasets_dir()
if not directory.is_dir():
return []
return sorted(directory.glob("*.json"))
def _content_hash(raw: bytes) -> str:
return "sha256:" + hashlib.sha256(raw).hexdigest()
def _read_json(path: Path) -> tuple[dict, bytes]:
"""读取并解析 JSON 文件,返回 (dict, 原始字节);非法 JSON 抛 BENCHMARK_DATASET_INVALID。"""
try:
raw_bytes = path.read_bytes()
return json.loads(raw_bytes.decode("utf-8")), raw_bytes
except (json.JSONDecodeError, OSError, UnicodeDecodeError) as exc:
raise ApiError(
422,
"BENCHMARK_DATASET_INVALID",
f"Dataset file is not valid JSON: {path.name}",
{"path": str(path)},
) from exc
def _dataset_from_raw(raw: dict, raw_bytes: bytes, kind: BenchmarkKind) -> RAGDataset:
"""把单个数据集 JSON 解析为 RAGDataset,非法结构抛 BENCHMARK_DATASET_INVALID。"""
dataset_id = raw.get("dataset_id")
if not isinstance(dataset_id, str) or not dataset_id:
raise ApiError(
422,
"BENCHMARK_DATASET_INVALID",
"Dataset must declare a non-empty string 'dataset_id'.",
{},
)
file_kind = raw.get("kind", kind.value)
if file_kind != kind.value:
raise ApiError(
422,
"BENCHMARK_DATASET_INVALID",
f"Dataset kind mismatch: expected '{kind.value}', got '{file_kind}'.",
{"dataset_id": dataset_id},
)
raw_cases = raw.get("cases")
if not isinstance(raw_cases, list) or not raw_cases:
raise ApiError(
422,
"BENCHMARK_DATASET_INVALID",
"Dataset 'cases' must be a non-empty list.",
{"dataset_id": dataset_id},
)
cases: list[RAGDatasetCase] = []
for index, case in enumerate(raw_cases):
try:
parsed = RAGDatasetCase.model_validate(case)
except ValidationError as exc:
raise ApiError(
422,
"BENCHMARK_DATASET_INVALID",
f"Dataset case #{index} is invalid.",
{"dataset_id": dataset_id, "case_index": index, "errors": exc.errors()},
) from exc
# 每个 Case 至少要声明一个期望 id,否则无法计算命中/召回
if not parsed.expected_note_ids and not parsed.expected_block_ids:
raise ApiError(
422,
"BENCHMARK_DATASET_INVALID",
f"Dataset case '{parsed.case_id}' must declare expected_note_ids or expected_block_ids.",
{"dataset_id": dataset_id, "case_id": parsed.case_id},
)
cases.append(parsed)
return RAGDataset(
dataset_id=dataset_id,
kind=kind,
version=str(raw.get("version", "")),
description=str(raw.get("description", "")),
cases=cases,
content_hash=_content_hash(raw_bytes),
)
def list_datasets(kind: BenchmarkKind) -> list[BenchmarkDatasetInfo]:
"""枚举受控目录下指定 kind 的数据集元信息(不含 Case 内容)。
个别文件损坏时跳过而非整体失败,保证列表接口健壮;损坏细节由 load_dataset 抛出。
"""
infos: list[BenchmarkDatasetInfo] = []
for path in _dataset_files():
try:
raw, raw_bytes = _read_json(path)
except ApiError:
continue
if raw.get("kind", kind.value) != kind.value:
continue
infos.append(
BenchmarkDatasetInfo(
dataset_id=raw.get("dataset_id", path.stem),
kind=kind,
version=str(raw.get("version", "")),
description=str(raw.get("description", "")),
case_count=len(raw.get("cases", [])),
content_hash=_content_hash(raw_bytes),
)
)
return infos
def load_dataset(dataset_id: str, kind: BenchmarkKind) -> RAGDataset:
"""按 id 加载并校验数据集;找不到抛 BENCHMARK_DATASET_NOT_FOUND。"""
for path in _dataset_files():
raw, raw_bytes = _read_json(path)
if raw.get("dataset_id") != dataset_id:
continue
return _dataset_from_raw(raw, raw_bytes, kind)
raise ApiError(
404,
"BENCHMARK_DATASET_NOT_FOUND",
f"Benchmark dataset does not exist: {dataset_id}",
{"dataset_id": dataset_id, "kind": kind.value},
)
+55
View File
@@ -0,0 +1,55 @@
"""Benchmark 指标纯函数。
所有指标只依赖「按相关性降序的 retrieved id 列表」和「期望 id 集合」,不接触任何
外部状态,便于单元测试与未来 Agent Benchmark 复用。retrieved 顺序越靠前越相关。
"""
from __future__ import annotations
def hit_at_k(retrieved: list[str], expected: set[str], k: int) -> bool:
"""前 k 个结果里是否命中任意期望 id(用于 Hit@1 / Hit@5)。"""
return any(item in expected for item in retrieved[:k])
def recall_at_k(retrieved: list[str], expected: set[str], k: int) -> float:
"""前 k 个结果召回的期望 id 占比;期望为空时视为 0。"""
if not expected:
return 0.0
hits = sum(1 for item in retrieved[:k] if item in expected)
return hits / len(expected)
def reciprocal_rank(retrieved: list[str], expected: set[str]) -> float:
"""首个命中的倒数排名;未命中返回 0。rank 从 1 开始。"""
for rank, item in enumerate(retrieved, start=1):
if item in expected:
return 1.0 / rank
return 0.0
def citation_hit(retrieved_block_ids: list[str], expected: set[str]) -> bool:
"""首条结果的 block_id 是否为期望引用块(Citation Hit Rate 的逐 Case 判据)。"""
if not retrieved_block_ids or not expected:
return False
return retrieved_block_ids[0] in expected
def mean(values: list[float]) -> float:
return sum(values) / len(values) if values else 0.0
def percentile(values: list[float], p: float) -> float:
"""线性插值分位数(p ∈ [0, 100]),用于 P50 / P95 延迟。空列表返回 0。"""
if not values:
return 0.0
ordered = sorted(values)
if len(ordered) == 1:
return ordered[0]
rank = (len(ordered) - 1) * (p / 100.0)
lo = int(rank)
hi = lo + 1
if hi >= len(ordered):
return ordered[-1]
frac = rank - lo
return ordered[lo] + (ordered[hi] - ordered[lo]) * frac
+113
View File
@@ -0,0 +1,113 @@
"""RAG Benchmark Runner:调用检索引擎对数据集逐 Case 求值并聚合指标。
只读操作,直接复用 app.retrieval.engine 的 search(),不旁路检索链路。指标按
(mode, case, repeat) 逐样本计算,再按 mode 聚合;失败样本保留在报告中但不计入汇总,
避免异常样本污染指标。
"""
from __future__ import annotations
import time
from collections.abc import Callable
from app.benchmarks import metrics as m
from app.benchmarks.datasets import RAGDataset
from app.contracts import (
RAGCaseResult,
RAGDatasetCase,
RAGMetrics,
RAGRunRequest,
SearchMode,
SearchRequest,
)
from app.retrieval.engine import engine
async def run_rag(
dataset: RAGDataset,
request: RAGRunRequest,
on_case: Callable[[RAGCaseResult, int, int], None] | None = None,
) -> tuple[dict[str, RAGMetrics], list[RAGCaseResult]]:
"""执行 RAG Benchmark,返回 (按 mode 聚合的指标, 全部逐样本结果)。
on_case 在每个样本求值完成后回调 (result, done, total),供上层更新进度与事件。
"""
total = len(request.modes) * len(dataset.cases) * request.repeat
done = 0
results: list[RAGCaseResult] = []
for mode in request.modes:
for case in dataset.cases:
for repeat in range(request.repeat):
result = await _evaluate_one(case, mode, request, repeat)
results.append(result)
done += 1
if on_case is not None:
on_case(result, done, total)
metrics_by_mode = {mode.value: _aggregate(results, mode) for mode in request.modes}
return metrics_by_mode, results
async def _evaluate_one(
case: RAGDatasetCase, mode: SearchMode, request: RAGRunRequest, repeat: int
) -> RAGCaseResult:
search_request = SearchRequest(
query=case.query,
mode=mode,
limit=request.retrieval.top_k,
include_snippet=False,
)
start = time.perf_counter()
try:
response = await engine.search(search_request)
latency_ms = (time.perf_counter() - start) * 1000.0
except Exception as exc: # 单个样本失败不中断整个 Benchmark
return RAGCaseResult(
case_id=case.case_id,
mode=mode,
repeat=repeat,
latency_ms=(time.perf_counter() - start) * 1000.0,
error=str(exc),
)
retrieved_note_ids = [item.note_id for item in response.items]
retrieved_block_ids = [item.block_id for item in response.items]
expected_notes = set(case.expected_note_ids)
expected_blocks = set(case.expected_block_ids)
k = request.retrieval.top_k
return RAGCaseResult(
case_id=case.case_id,
mode=mode,
repeat=repeat,
latency_ms=latency_ms,
retrieved_note_ids=retrieved_note_ids,
retrieved_block_ids=retrieved_block_ids,
hit_at_1=m.hit_at_k(retrieved_note_ids, expected_notes, 1),
hit_at_5=m.hit_at_k(retrieved_note_ids, expected_notes, 5),
recall=m.recall_at_k(retrieved_note_ids, expected_notes, k),
reciprocal_rank=m.reciprocal_rank(retrieved_note_ids, expected_notes),
citation_hit=m.citation_hit(retrieved_block_ids, expected_blocks),
citation_applicable=bool(case.expected_block_ids),
)
def _aggregate(cases: list[RAGCaseResult], mode: SearchMode) -> RAGMetrics:
samples = [c for c in cases if c.mode == mode]
ok = [c for c in samples if c.error is None]
if not ok:
return RAGMetrics()
latencies = [c.latency_ms for c in ok]
# citation_hit_rate 只统计声明了 expected_block_ids 的样本
citation_samples = [c for c in ok if c.citation_applicable]
return RAGMetrics(
hit_at_1=m.mean([1.0 if c.hit_at_1 else 0.0 for c in ok]),
hit_at_5=m.mean([1.0 if c.hit_at_5 else 0.0 for c in ok]),
recall_at_k=m.mean([c.recall for c in ok]),
mrr=m.mean([c.reciprocal_rank for c in ok]),
citation_hit_rate=m.mean([1.0 if c.citation_hit else 0.0 for c in citation_samples]),
p50_latency_ms=m.percentile(latencies, 50.0),
p95_latency_ms=m.percentile(latencies, 95.0),
)
+193
View File
@@ -0,0 +1,193 @@
"""Benchmark 服务:运行注册表、配置快照与报告组装。
MVP 阶段运行是同步的(与 index_service 一致):POST 创建后立即执行完并返回
completed 的 BenchmarkRun。运行记录、事件与报告暂存内存(_runs/_events/_reports),
不持久化到 SQLite;后续接入异步任务队列时再落库。
"""
from __future__ import annotations
import sys
from datetime import datetime, timezone
from uuid import uuid4
from app.benchmarks import datasets
from app.benchmarks.datasets import RAGDataset
from app.benchmarks.rag import run_rag
from app.config import get_settings
from app.contracts import (
BenchmarkEvent,
BenchmarkEventType,
BenchmarkKind,
BenchmarkReport,
BenchmarkRun,
BenchmarkStatus,
RAGCaseResult,
RAGMetrics,
RAGRunRequest,
)
from app.errors import ApiError
from app.retrieval.engine import engine
_runs: dict[str, BenchmarkRun] = {}
_events: dict[str, list[BenchmarkEvent]] = {}
_reports: dict[str, BenchmarkReport] = {}
MAX_RUNS = 100
def _now() -> datetime:
return datetime.now(timezone.utc)
def _remember(run: BenchmarkRun) -> None:
_runs[run.run_id] = run
while len(_runs) > MAX_RUNS:
oldest = next(iter(_runs))
_runs.pop(oldest, None)
_events.pop(oldest, None)
_reports.pop(oldest, None)
def _config_snapshot(request: RAGRunRequest, dataset: RAGDataset) -> dict:
"""记录运行时的模型 / 索引 / 环境信息,保证报告可解释、可复现。"""
settings = get_settings()
return {
"dataset_id": dataset.dataset_id,
"dataset_hash": dataset.content_hash,
"dataset_version": dataset.version,
"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},
"app": {"version": settings.version, "environment": settings.environment},
"python": sys.version.split()[0],
"metadata": request.metadata,
}
async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
"""创建并同步执行一次 RAG Benchmark,返回 completed 的 BenchmarkRun。"""
dataset = datasets.load_dataset(request.dataset_id, BenchmarkKind.rag)
run_id = "benchmark_" + uuid4().hex[:12]
snapshot = _config_snapshot(request, dataset)
run = BenchmarkRun(
run_id=run_id,
kind=BenchmarkKind.rag,
dataset_id=dataset.dataset_id,
dataset_hash=dataset.content_hash,
status=BenchmarkStatus.running,
progress=0.0,
config_snapshot=snapshot,
created_at=_now(),
started_at=_now(),
)
_remember(run)
_events[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(),
)
)
emit(
BenchmarkEventType.run_started,
{"dataset_id": dataset.dataset_id, "modes": [m.value for m in request.modes]},
)
total = len(request.modes) * len(dataset.cases) * request.repeat
def on_case(result: RAGCaseResult, done: int, _total: int) -> None:
progress = done / total if total else 1.0
_runs[run_id] = _runs[run_id].model_copy(update={"progress": progress})
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:
_runs[run_id] = _runs[run_id].model_copy(
update={
"status": BenchmarkStatus.failed,
"progress": 1.0,
"error": str(exc),
"completed_at": _now(),
}
)
emit(BenchmarkEventType.run_failed, {"error": str(exc)})
_reports[run_id] = BenchmarkReport(
run_id=run_id,
kind=BenchmarkKind.rag,
dataset_id=dataset.dataset_id,
dataset_hash=dataset.content_hash,
status=BenchmarkStatus.failed,
config_snapshot=snapshot,
error=str(exc),
)
raise ApiError(500, "BENCHMARK_RUN_FAILED", str(exc), {"run_id": run_id}) from exc
metrics = {mode: m.model_dump() for mode, m in metrics_by_mode.items()}
_runs[run_id] = _runs[run_id].model_copy(
update={
"status": BenchmarkStatus.completed,
"progress": 1.0,
"metrics": metrics,
"completed_at": _now(),
}
)
emit(BenchmarkEventType.run_completed, {"metrics": metrics})
_reports[run_id] = BenchmarkReport(
run_id=run_id,
kind=BenchmarkKind.rag,
dataset_id=dataset.dataset_id,
dataset_hash=dataset.content_hash,
status=BenchmarkStatus.completed,
config_snapshot=snapshot,
metrics=metrics,
cases=results,
)
return _runs[run_id]
def list_runs(
kind: BenchmarkKind | None = None,
status: BenchmarkStatus | None = None,
limit: int = 50,
offset: int = 0,
) -> tuple[list[BenchmarkRun], int]:
runs = list(_runs.values())
if kind is not None:
runs = [r for r in runs if r.kind == kind]
if status is not None:
runs = [r for r in runs if r.status == status]
runs.sort(key=lambda r: r.created_at, reverse=True)
total = len(runs)
return runs[offset : offset + limit], total
def get_run(run_id: str) -> BenchmarkRun | None:
return _runs.get(run_id)
def get_report(run_id: str) -> BenchmarkReport | None:
return _reports.get(run_id)
def get_events(run_id: str) -> list[BenchmarkEvent]:
return _events.get(run_id, [])
def cancel_run(run_id: str) -> BenchmarkRun | None:
"""取消运行:同步 MVP 下运行通常已结束,仅对仍在排队/运行的记录置为 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
return run
+4
View File
@@ -24,6 +24,7 @@ class Settings:
db_path: Path db_path: Path
vault_path: Path vault_path: Path
attachments_path: Path attachments_path: Path
benchmark_datasets_path: Path
@lru_cache @lru_cache
@@ -41,4 +42,7 @@ def get_settings() -> Settings:
attachments_path=Path( attachments_path=Path(
os.getenv("APP_ATTACHMENTS_PATH", str(data_dir / "attachments")) os.getenv("APP_ATTACHMENTS_PATH", str(data_dir / "attachments"))
), ),
benchmark_datasets_path=Path(
os.getenv("APP_BENCHMARK_DATASETS_PATH", str(data_dir / "benchmarks"))
),
) )
+131
View File
@@ -626,3 +626,134 @@ class IndexJob(Contract):
status: Literal["queued", "running", "completed", "failed"] status: Literal["queued", "running", "completed", "failed"]
scope: Literal["all", "notes", "vectors"] scope: Literal["all", "notes", "vectors"]
created_at: datetime created_at: datetime
# Benchmark
class BenchmarkKind(str, Enum):
rag = "rag"
agent = "agent"
class BenchmarkStatus(str, Enum):
queued = "queued"
running = "running"
completed = "completed"
failed = "failed"
cancelled = "cancelled"
class RAGDatasetCase(Contract):
case_id: str
query: str = Field(min_length=1)
expected_note_ids: list[str] = Field(default_factory=list)
expected_block_ids: list[str] = Field(default_factory=list)
citation_required: bool = False
tags: list[str] = Field(default_factory=list)
class RAGRetrievalConfig(Contract):
"""RAG Benchmark 的检索参数。top_k 映射到 SearchRequest.limit,其余参数当前
记录进 config_snapshot,由 RetrievalProfile 共享(§9.9)落地后再接入引擎。"""
top_k: int = Field(default=10, ge=1, le=100)
rrf_k: int = Field(default=60, ge=1)
rerank: bool = True
rerank_candidates: int = Field(default=20, ge=1)
score_threshold: float = Field(default=0.0, ge=0.0)
class RAGRunRequest(Contract):
dataset_id: str = Field(min_length=1)
modes: list[SearchMode] = Field(
default_factory=lambda: [SearchMode.fts, SearchMode.vector, SearchMode.hybrid]
)
retrieval: RAGRetrievalConfig = Field(default_factory=RAGRetrievalConfig)
repeat: int = Field(default=1, ge=1, le=10)
metadata: dict[str, Any] = Field(default_factory=dict)
class RAGMetrics(Contract):
hit_at_1: float = 0.0
hit_at_5: float = 0.0
recall_at_k: float = 0.0
mrr: float = 0.0
citation_hit_rate: float = 0.0
p50_latency_ms: float = 0.0
p95_latency_ms: float = 0.0
class BenchmarkDatasetInfo(Contract):
dataset_id: str
kind: BenchmarkKind
version: str
description: str = ""
case_count: int
content_hash: str
class BenchmarkDatasetListResponse(Contract):
items: list[BenchmarkDatasetInfo] = Field(default_factory=list)
class BenchmarkRun(Contract):
run_id: str
kind: BenchmarkKind
dataset_id: str
dataset_hash: str
status: BenchmarkStatus
progress: float | None = None
metrics: dict[str, Any] | None = None
config_snapshot: dict[str, Any] = Field(default_factory=dict)
error: str | None = None
created_at: datetime
started_at: datetime | None = None
completed_at: datetime | None = None
class BenchmarkRunListResponse(Contract):
items: list[BenchmarkRun] = Field(default_factory=list)
page: PageMeta = Field(default_factory=PageMeta)
class BenchmarkEventType(str, Enum):
run_started = "RunStarted"
case_completed = "CaseCompleted"
run_completed = "RunCompleted"
run_failed = "RunFailed"
class BenchmarkEvent(Contract):
event: BenchmarkEventType
run_id: str
sequence: int
data: dict[str, Any] = Field(default_factory=dict)
timestamp: datetime
class RAGCaseResult(Contract):
case_id: str
mode: SearchMode
repeat: int
latency_ms: float
retrieved_note_ids: list[str] = Field(default_factory=list)
retrieved_block_ids: list[str] = Field(default_factory=list)
hit_at_1: bool = False
hit_at_5: bool = False
recall: float = 0.0
reciprocal_rank: float = 0.0
citation_hit: bool = False
# 该 Case 是否声明了 expected_block_ids(决定是否计入 citation_hit_rate 分母)
citation_applicable: bool = False
error: str | None = None
class BenchmarkReport(Contract):
run_id: str
kind: BenchmarkKind
dataset_id: str
dataset_hash: str
status: BenchmarkStatus
config_snapshot: dict[str, Any] = Field(default_factory=dict)
metrics: dict[str, Any] = Field(default_factory=dict)
cases: list[RAGCaseResult] = Field(default_factory=list)
error: str | None = None
+138 -1
View File
@@ -11,6 +11,13 @@ from app.contracts import (
AgentRunListResponse, AgentRunListResponse,
AgentTraceResponse, AgentTraceResponse,
ChatRequest, ChatRequest,
BenchmarkDatasetListResponse,
BenchmarkKind,
BenchmarkReport,
BenchmarkRun,
BenchmarkRunListResponse,
BenchmarkStatus,
RAGRunRequest,
CredentialStatus, CredentialStatus,
CredentialWriteRequest, CredentialWriteRequest,
ExtensionInstallRequest, ExtensionInstallRequest,
@@ -59,8 +66,10 @@ from app.contracts import (
WorkspaceSnapshot, WorkspaceSnapshot,
) )
from app.agent import AgentCapacityError, AgentRunNotFoundError from app.agent import AgentCapacityError, AgentRunNotFoundError
from app.benchmarks import datasets as benchmark_datasets
from app.benchmarks import service as benchmark_service
from app.container import container from app.container import container
from app.errors import ApiError from app.errors import ApiError, not_implemented
from app.extensions import ExtensionError from app.extensions import ExtensionError
from app.providers.registry import ProviderNotFoundError from app.providers.registry import ProviderNotFoundError
from app.providers.factory import UnsupportedProviderError from app.providers.factory import UnsupportedProviderError
@@ -795,3 +804,131 @@ async def get_index_job(job_id: str) -> IndexJob:
if job is None: if job is None:
raise ApiError(404, "RESOURCE_NOT_FOUND", "index job not found", {"job_id": job_id}) raise ApiError(404, "RESOURCE_NOT_FOUND", "index job not found", {"job_id": job_id})
return job return job
# Benchmark
@router.get(
"/benchmarks/datasets",
response_model=BenchmarkDatasetListResponse,
tags=["Benchmark"],
)
async def list_benchmark_datasets(
kind: BenchmarkKind = Query(default=BenchmarkKind.rag),
) -> BenchmarkDatasetListResponse:
return BenchmarkDatasetListResponse(items=benchmark_datasets.list_datasets(kind))
@router.post(
"/benchmarks/rag/runs",
response_model=BenchmarkRun,
status_code=202,
tags=["Benchmark"],
)
async def create_rag_benchmark(request: RAGRunRequest) -> BenchmarkRun:
return await benchmark_service.create_rag_run(request)
@router.post(
"/benchmarks/agent/runs",
response_model=BenchmarkRun,
status_code=202,
tags=["Benchmark"],
)
async def create_agent_benchmark() -> BenchmarkRun:
# Agent Benchmark 基础设施在 RAG Benchmark 之后单独交付,先占位契约
not_implemented("Agent Benchmark")
raise AssertionError("unreachable")
@router.get(
"/benchmarks/runs",
response_model=BenchmarkRunListResponse,
tags=["Benchmark"],
)
async def list_benchmark_runs(
kind: BenchmarkKind | None = Query(default=None),
status: BenchmarkStatus | None = Query(default=None),
limit: int = Query(default=50, ge=1, le=100),
offset: int = Query(default=0, ge=0),
) -> BenchmarkRunListResponse:
items, total = benchmark_service.list_runs(
kind=kind, status=status, limit=limit, offset=offset
)
return BenchmarkRunListResponse(
items=items, page=PageMeta(total=total, limit=limit, offset=offset)
)
@router.get(
"/benchmarks/runs/{run_id}",
response_model=BenchmarkRun,
tags=["Benchmark"],
)
async def get_benchmark_run(run_id: str) -> BenchmarkRun:
run = benchmark_service.get_run(run_id)
if run is None:
raise ApiError(
404, "BENCHMARK_RUN_NOT_FOUND", "benchmark run not found", {"run_id": run_id}
)
return run
@router.post(
"/benchmarks/runs/{run_id}/cancel",
response_model=OperationResponse,
tags=["Benchmark"],
)
async def cancel_benchmark_run(run_id: str) -> OperationResponse:
run = benchmark_service.cancel_run(run_id)
if run is None:
raise ApiError(
404, "BENCHMARK_RUN_NOT_FOUND", "benchmark run not found", {"run_id": run_id}
)
return OperationResponse(
status="completed",
resource_id=run_id,
message=f"Benchmark run status: {run.status.value}",
)
@router.get(
"/benchmarks/runs/{run_id}/events",
response_class=StreamingResponse,
responses={
200: {
"description": "BenchmarkEvent Server-Sent Events stream",
"content": {"text/event-stream": {}},
}
},
tags=["Benchmark"],
)
async def benchmark_events(
run_id: str,
after_sequence: int = Query(default=-1, ge=-1),
) -> StreamingResponse:
if benchmark_service.get_run(run_id) is None:
raise ApiError(
404, "BENCHMARK_RUN_NOT_FOUND", "benchmark run not found", {"run_id": run_id}
)
async def stream() -> AsyncIterator[str]:
for event in benchmark_service.get_events(run_id):
if event.sequence <= after_sequence:
continue
yield as_sse(event.event.value, event.model_dump_json(), event_id=event.sequence)
return StreamingResponse(stream(), media_type="text/event-stream")
@router.get(
"/benchmarks/runs/{run_id}/report",
response_model=BenchmarkReport,
tags=["Benchmark"],
)
async def get_benchmark_report(run_id: str) -> BenchmarkReport:
report = benchmark_service.get_report(run_id)
if report is None:
raise ApiError(
404, "BENCHMARK_RUN_NOT_FOUND", "benchmark report not found", {"run_id": run_id}
)
return report
+48
View File
@@ -0,0 +1,48 @@
{
"dataset_id": "rag-core-v1",
"kind": "rag",
"version": "1.0.0",
"description": "基础中文笔记检索集(对应 backend/data/vault 内置语料,重建索引后即可复现)",
"cases": [
{
"case_id": "rag-vector-sim",
"query": "向量数据库如何进行相似度检索",
"expected_note_ids": ["note_c1454740a0e55ef5"],
"expected_block_ids": ["blk_07c4c6bce0ec4d12", "blk_605fb3593809f224"],
"citation_required": true,
"tags": ["向量数据库", "检索"]
},
{
"case_id": "rag-python-func",
"query": "Python 如何定义函数",
"expected_note_ids": ["note_424c3742c6f0e555"],
"expected_block_ids": ["blk_45d48cae2fed40fe", "blk_0768d9c25c2ecf07"],
"citation_required": true,
"tags": ["python"]
},
{
"case_id": "rag-citation",
"query": "搜索结果如何定位到原文位置",
"expected_note_ids": ["note_0c619caa30b1614c"],
"expected_block_ids": ["blk_3f6fcead71c25fc6", "blk_9af7b12e9ce909fc"],
"citation_required": true,
"tags": ["RAG"]
},
{
"case_id": "rag-hybrid",
"query": "混合检索怎么融合全文和向量",
"expected_note_ids": ["note_c1454740a0e55ef5"],
"expected_block_ids": ["blk_82b45418dba9f720"],
"citation_required": true,
"tags": ["检索"]
},
{
"case_id": "rag-tech-stack",
"query": "这个项目用什么后端和检索技术",
"expected_note_ids": ["note_3327e6cf18f3701f"],
"expected_block_ids": ["blk_feb2a9c42e7d31ad"],
"citation_required": false,
"tags": ["项目"]
}
]
}
@@ -0,0 +1,3 @@
# 第一阶段验收
Notes Agent 支持混合检索和可定位引用。
+205
View File
@@ -0,0 +1,205 @@
"""Benchmark 服务的单元与端到端测试。
沿用 conftest 的隔离机制:APP_DATA_DIR / DB / Vault 都指向临时目录,benchmark
数据集也落在临时目录(settings.benchmark_datasets_path),不读写真实数据。
"""
from __future__ import annotations
import asyncio
import json
import pytest
from app.benchmarks import datasets, metrics as m, service
from app.config import get_settings
from app.contracts import BenchmarkKind, RAGRunRequest, SearchMode
from app.errors import ApiError
def _write_dataset(dataset_id: str, cases: list[dict], *, kind: str = "rag") -> None:
directory = get_settings().benchmark_datasets_path
directory.mkdir(parents=True, exist_ok=True)
payload = {
"dataset_id": dataset_id,
"kind": kind,
"version": "1.0.0",
"description": "test dataset",
"cases": cases,
}
(directory / f"{dataset_id}.json").write_text(
json.dumps(payload, ensure_ascii=False), encoding="utf-8"
)
# --------------------------------------------------------------------------- #
# 指标纯函数
# --------------------------------------------------------------------------- #
def test_hit_at_k_and_recall() -> None:
retrieved = ["a", "b", "c"]
expected = {"b", "z"}
assert m.hit_at_k(retrieved, expected, 1) is False
assert m.hit_at_k(retrieved, expected, 2) is True
assert m.recall_at_k(retrieved, expected, 5) == 0.5 # 只召回 b
def test_reciprocal_rank_and_citation_hit() -> None:
assert m.reciprocal_rank(["x", "a", "b"], {"b"}) == 1 / 3
assert m.reciprocal_rank(["x"], {"b"}) == 0.0
assert m.citation_hit(["blk_1"], {"blk_1"}) is True
assert m.citation_hit(["blk_2"], {"blk_1"}) is False
assert m.citation_hit([], {"blk_1"}) is False
def test_percentile() -> None:
assert m.percentile([1.0, 2.0, 3.0, 4.0], 50.0) == 2.5
assert m.percentile([], 50.0) == 0.0
assert m.percentile([7.0], 95.0) == 7.0
# --------------------------------------------------------------------------- #
# Dataset 注册与校验
# --------------------------------------------------------------------------- #
def test_list_datasets_empty_by_default() -> None:
assert datasets.list_datasets(BenchmarkKind.rag) == []
def test_load_missing_dataset_raises() -> None:
with pytest.raises(ApiError) as exc:
datasets.load_dataset("does-not-exist", BenchmarkKind.rag)
assert exc.value.status_code == 404
assert exc.value.code == "BENCHMARK_DATASET_NOT_FOUND"
def test_dataset_without_expected_ids_is_invalid() -> None:
_write_dataset("bad-v1", [{"case_id": "x", "query": "q", "citation_required": False}])
with pytest.raises(ApiError) as exc:
datasets.load_dataset("bad-v1", BenchmarkKind.rag)
assert exc.value.code == "BENCHMARK_DATASET_INVALID"
def test_dataset_kind_mismatch_is_invalid() -> None:
_write_dataset("agent-v1", [{"case_id": "x", "query": "q", "expected_note_ids": ["n"]}], kind="agent")
with pytest.raises(ApiError) as exc:
datasets.load_dataset("agent-v1", BenchmarkKind.rag)
assert exc.value.code == "BENCHMARK_DATASET_INVALID"
# --------------------------------------------------------------------------- #
# RAG Benchmark 端到端
# --------------------------------------------------------------------------- #
def _single_note_case() -> tuple[str, str, dict]:
from app.services import note_service
note = asyncio.run(
note_service.create_note(
title="向量库",
markdown="向量数据库用于存储高维向量并支持近似最近邻检索。",
folder="",
tags=["向量"],
)
)
case = {
"case_id": "c1",
"query": "向量数据库相似度检索",
"expected_note_ids": [note.note_id],
"expected_block_ids": [note.blocks[0].block_id],
"citation_required": True,
"tags": ["向量"],
}
return note.note_id, note.blocks[0].block_id, case
def test_rag_benchmark_end_to_end() -> None:
_, _, case = _single_note_case()
_write_dataset("e2e-v1", [case])
run = asyncio.run(
service.create_rag_run(RAGRunRequest(dataset_id="e2e-v1", modes=[SearchMode.fts]))
)
assert run.status.value == "completed"
assert run.dataset_hash.startswith("sha256:")
assert run.metrics is not None
fts = run.metrics["fts"]
assert fts["hit_at_1"] == 1.0
assert fts["recall_at_k"] == 1.0
assert fts["mrr"] == 1.0
assert fts["citation_hit_rate"] == 1.0
assert fts["p50_latency_ms"] >= 0.0
assert fts["p95_latency_ms"] >= fts["p50_latency_ms"]
def test_rag_benchmark_all_modes_produce_metrics() -> None:
_, _, case = _single_note_case()
_write_dataset("e2e-modes-v1", [case])
run = asyncio.run(service.create_rag_run(RAGRunRequest(dataset_id="e2e-modes-v1")))
assert run.status.value == "completed"
for mode in ("fts", "vector", "hybrid"):
assert mode in run.metrics
for key in ("hit_at_1", "hit_at_5", "recall_at_k", "mrr", "citation_hit_rate"):
assert 0.0 <= run.metrics[mode][key] <= 1.0
def test_benchmark_report_and_events() -> None:
_, _, case = _single_note_case()
_write_dataset("report-v1", [case])
run = asyncio.run(service.create_rag_run(RAGRunRequest(dataset_id="report-v1", modes=[SearchMode.fts])))
report = service.get_report(run.run_id)
events = service.get_events(run.run_id)
assert report is not None
assert report.run_id == run.run_id
assert len(report.cases) == 1
assert report.cases[0].case_id == "c1"
assert report.cases[0].hit_at_1 is True
assert events, "运行应产生事件"
assert events[0].event.value == "RunStarted"
assert events[-1].event.value == "RunCompleted"
def test_cancel_completed_run_keeps_status() -> None:
_, _, case = _single_note_case()
_write_dataset("cancel-v1", [case])
run = asyncio.run(service.create_rag_run(RAGRunRequest(dataset_id="cancel-v1", modes=[SearchMode.fts])))
cancelled = service.cancel_run(run.run_id)
assert cancelled.status.value == "completed" # 同步运行已结束,不再变 cancelled
# --------------------------------------------------------------------------- #
# 路由接入
# --------------------------------------------------------------------------- #
def test_benchmark_routes_wired() -> None:
from app import routes
_, _, case = _single_note_case()
_write_dataset("route-v1", [case])
listed = asyncio.run(routes.list_benchmark_datasets(BenchmarkKind.rag))
assert any(item.dataset_id == "route-v1" for item in listed.items)
run = asyncio.run(
routes.create_rag_benchmark(RAGRunRequest(dataset_id="route-v1", modes=[SearchMode.fts]))
)
assert run.status.value == "completed"
got = asyncio.run(routes.get_benchmark_run(run.run_id))
assert got.run_id == run.run_id
report = asyncio.run(routes.get_benchmark_report(run.run_id))
assert report.cases[0].case_id == "c1"
def test_benchmark_run_not_found_raises() -> None:
from app import routes
with pytest.raises(ApiError) as exc:
asyncio.run(routes.get_benchmark_run("benchmark_missing"))
assert exc.value.code == "BENCHMARK_RUN_NOT_FOUND"