feat: 完成benchmark后端功能
This commit is contained in:
@@ -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 运行注册表、配置快照与报告组装
|
||||||
|
"""
|
||||||
@@ -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},
|
||||||
|
)
|
||||||
@@ -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
|
||||||
@@ -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),
|
||||||
|
)
|
||||||
@@ -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
|
||||||
@@ -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"))
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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 支持混合检索和可定位引用。
|
||||||
@@ -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"
|
||||||
Reference in New Issue
Block a user