Files
NotesAgentic/backend/tests/test_benchmark.py
T

568 lines
21 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.
"""Benchmark 服务的单元与端到端测试。
沿用 conftest 的隔离机制:APP_DATA_DIR / DB / Vault 都指向临时目录,benchmark
数据集也落在临时目录(settings.benchmark_datasets_path),不读写真实数据。
运行采用「创建即 queued + 后台 Task 执行」的异步模型,测试通过 _run 在同一事件循环内
创建并等待后台任务结束,得到终态 BenchmarkRun 后再断言。
"""
from __future__ import annotations
import asyncio
import json
import pytest
from pydantic import ValidationError
from app.benchmarks import datasets, metrics as m, service
from app.config import get_settings
from app.contracts import (
BenchmarkKind,
BenchmarkRun,
BenchmarkStatus,
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 _write_raw(dataset_id: str, raw: dict) -> None:
directory = get_settings().benchmark_datasets_path
directory.mkdir(parents=True, exist_ok=True)
(directory / f"{dataset_id}.json").write_text(
json.dumps(raw, ensure_ascii=False), encoding="utf-8"
)
def _run(request: RAGRunRequest):
"""创建运行并在同一事件循环内等待后台任务结束,返回终态 BenchmarkRun。"""
from app.contracts import BenchmarkRun
async def _execute() -> BenchmarkRun:
run = await service.create_rag_run(request)
return await service.wait_for_run(run.run_id)
return asyncio.run(_execute())
# --------------------------------------------------------------------------- #
# 指标纯函数
# --------------------------------------------------------------------------- #
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_recall_at_k_dedups_duplicate_notes() -> None:
# 同一 Note 经多个 Block 重复出现,去重后 Recall 不应超过 1
assert m.recall_at_k(["note-a", "note-a"], {"note-a"}, 2) == 1.0
assert m.recall_at_k(["note-a", "note-a", "note-b"], {"note-a"}, 3) == 1.0
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"
def test_citation_required_requires_expected_block_ids() -> None:
# citation_required=true 却没有 expected_block_ids,无法计算 Citation Hit Rate,应拒绝
_write_dataset(
"cit-req-v1",
[{"case_id": "x", "query": "q", "expected_note_ids": ["n"], "citation_required": True}],
)
with pytest.raises(ApiError) as exc:
datasets.load_dataset("cit-req-v1", BenchmarkKind.rag)
assert exc.value.code == "BENCHMARK_DATASET_INVALID"
def test_list_datasets_skips_corrupted_structure() -> None:
# 合法 JSON 但字段结构错误(cases: 42),列表接口应隔离该文件而非整体 500
_write_raw("bad-structure", {"dataset_id": "bad-structure", "kind": "rag", "cases": 42})
_write_dataset("good-v1", [{"case_id": "x", "query": "q", "expected_note_ids": ["n"]}])
infos = datasets.list_datasets(BenchmarkKind.rag)
ids = {info.dataset_id for info in infos}
assert "good-v1" in ids
assert "bad-structure" not in ids
# --------------------------------------------------------------------------- #
# 请求校验(空 / 重复 modes)
# --------------------------------------------------------------------------- #
def test_empty_modes_rejected() -> None:
with pytest.raises(ValidationError):
RAGRunRequest(dataset_id="x", modes=[])
def test_duplicate_modes_rejected() -> None:
with pytest.raises(ValidationError):
RAGRunRequest(dataset_id="x", modes=[SearchMode.fts, SearchMode.fts])
# --------------------------------------------------------------------------- #
# 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 = _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 = _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_config_snapshot_records_index_and_models() -> None:
_, _, case = _single_note_case()
_write_dataset("snapshot-v1", [case])
run = _run(RAGRunRequest(dataset_id="snapshot-v1", modes=[SearchMode.fts]))
snapshot = run.config_snapshot
assert snapshot["index_meta"] is not None
assert snapshot["embedding"]["policy"] == "per_case"
assert snapshot["local_embedding"]["version"]
assert snapshot["local_embedding"]["dim"]
assert snapshot["reranker"]["version"]
assert snapshot["retrieval"]["rrf_k"] == 60
def test_benchmark_report_and_events() -> None:
_, _, case = _single_note_case()
_write_dataset("report-v1", [case])
run = _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 = _run(RAGRunRequest(dataset_id="cancel-v1", modes=[SearchMode.fts]))
assert run.status.value == "completed"
cancelled = service.cancel_run(run.run_id)
assert cancelled.status.value == "completed" # 已结束,不再变 cancelled
def test_cancel_queued_run_marks_cancelled() -> None:
_, _, case = _single_note_case()
_write_dataset("cancel-queued-v1", [case])
async def _scenario():
run = await service.create_rag_run(
RAGRunRequest(dataset_id="cancel-queued-v1", modes=[SearchMode.fts])
)
service.cancel_run(run.run_id)
return await service.wait_for_run(run.run_id)
run = asyncio.run(_scenario())
assert run.status.value == "cancelled"
# --------------------------------------------------------------------------- #
# 指标聚合:Citation Hit Rate 只统计 citation_required 样本
# --------------------------------------------------------------------------- #
def test_citation_hit_rate_only_counts_citation_required() -> None:
from app.benchmarks import rag as rag_module
from app.contracts import RAGCaseResult
cases = [
RAGCaseResult(
case_id="a", mode=SearchMode.fts, repeat=0, latency_ms=1.0,
citation_hit=True, citation_applicable=True,
),
RAGCaseResult(
case_id="b", mode=SearchMode.fts, repeat=0, latency_ms=1.0,
citation_hit=False, citation_applicable=False,
),
]
metrics = rag_module._aggregate(cases, SearchMode.fts)
# 只有 citation_applicablecitation_required=true)的样本计入分母
assert metrics.citation_hit_rate == 1.0
# --------------------------------------------------------------------------- #
# 路由接入
# --------------------------------------------------------------------------- #
def test_benchmark_routes_wired() -> None:
from app import routes
_, _, case = _single_note_case()
_write_dataset("route-v1", [case])
async def _scenario():
listed = await routes.list_benchmark_datasets(BenchmarkKind.rag)
assert any(item.dataset_id == "route-v1" for item in listed.items)
run = await routes.create_rag_benchmark(
RAGRunRequest(dataset_id="route-v1", modes=[SearchMode.fts])
)
assert run.status.value == "queued"
return await service.wait_for_run(run.run_id)
run = asyncio.run(_scenario())
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"
# --------------------------------------------------------------------------- #
# 审阅回归:索引兼容 / 容量 / 失败样本 / 取消事件 / 数据集隔离
# --------------------------------------------------------------------------- #
def test_create_rag_run_requires_built_index() -> None:
# 空索引(无已索引 block)会让所有模式得到全 0 指标,应在创建时拒绝而非跑出误导结果
_write_dataset("empty-index-v1", [{"case_id": "x", "query": "q", "expected_note_ids": ["n"]}])
with pytest.raises(ApiError) as exc:
asyncio.run(
service.create_rag_run(
RAGRunRequest(dataset_id="empty-index-v1", modes=[SearchMode.fts])
)
)
assert exc.value.status_code == 409
assert exc.value.code == "BENCHMARK_INDEX_INCOMPATIBLE"
def test_capacity_exceeded_when_all_runs_active(monkeypatch) -> None:
# 满容量且全为活动(非终态)run 时,无法淘汰,应拒绝创建而非删掉正在运行的 run
_, _, case = _single_note_case()
_write_dataset("capacity-v1", [case])
monkeypatch.setattr(service, "MAX_RUNS", 1)
fake_id = "benchmark_fake_active"
service._runs[fake_id] = BenchmarkRun(
run_id=fake_id,
kind=BenchmarkKind.rag,
dataset_id="capacity-v1",
dataset_hash="sha256:fake",
status=BenchmarkStatus.queued,
created_at=service._now(),
)
try:
with pytest.raises(ApiError) as exc:
asyncio.run(
service.create_rag_run(
RAGRunRequest(dataset_id="capacity-v1", modes=[SearchMode.fts])
)
)
assert exc.value.status_code == 429
assert exc.value.code == "BENCHMARK_CAPACITY_EXCEEDED"
finally:
service._runs.pop(fake_id, None)
def test_failed_samples_counted_as_zero_in_aggregate() -> None:
from app.benchmarks import rag as rag_module
from app.contracts import RAGCaseResult
cases = [
RAGCaseResult(
case_id="ok", mode=SearchMode.fts, repeat=0, latency_ms=10.0,
hit_at_1=True, recall=1.0, reciprocal_rank=1.0,
citation_hit=True, citation_applicable=True,
),
RAGCaseResult(
case_id="boom", mode=SearchMode.fts, repeat=0, latency_ms=0.0,
error="RAG case evaluation failed.",
error_code="BENCHMARK_CASE_EVALUATION_FAILED",
),
]
metrics = rag_module._aggregate(cases, SearchMode.fts)
assert metrics.total_cases == 2
assert metrics.successful_cases == 1
assert metrics.failed_cases == 1
assert metrics.failure_rate == 0.5
# 失败样本按零分计入质量指标分母,汇总不虚高
assert metrics.hit_at_1 == 0.5
assert metrics.recall_at_k == 0.5
# 延迟只统计成功样本
assert metrics.p50_latency_ms == 10.0
def test_cancel_emits_run_cancelled_event() -> None:
_, _, case = _single_note_case()
_write_dataset("cancel-event-v1", [case])
async def _scenario():
run = await service.create_rag_run(
RAGRunRequest(dataset_id="cancel-event-v1", modes=[SearchMode.fts])
)
service.cancel_run(run.run_id)
return await service.wait_for_run(run.run_id)
run = asyncio.run(_scenario())
assert run.status.value == "cancelled"
events = service.get_events(run.run_id)
assert events[-1].event.value == "RunCancelled"
def test_load_dataset_ignores_corrupted_unrelated_files() -> None:
# 无关文件损坏(非法 JSON / 顶层非对象)不应阻断目标数据集加载
directory = get_settings().benchmark_datasets_path
directory.mkdir(parents=True, exist_ok=True)
(directory / "broken.json").write_text("{ not valid json", encoding="utf-8")
(directory / "array.json").write_text('["a", "b"]', encoding="utf-8")
_write_dataset("ok-v1", [{"case_id": "x", "query": "q", "expected_note_ids": ["n"]}])
dataset = datasets.load_dataset("ok-v1", BenchmarkKind.rag)
assert dataset.dataset_id == "ok-v1"
assert len(dataset.cases) == 1
def test_load_dataset_top_level_must_be_object() -> None:
_write_raw("array-top", ["a", "b"])
with pytest.raises(ApiError) as exc:
datasets.load_dataset("array-top", BenchmarkKind.rag)
assert exc.value.code == "BENCHMARK_DATASET_INVALID"
# --------------------------------------------------------------------------- #
# 审阅回归:运行中取消 / 仅块标注 / SSE 终止事件
# --------------------------------------------------------------------------- #
def test_cancel_running_benchmark_stops_early() -> None:
"""运行中取消应在样本边界及时生效,而非跑完全部样本(审阅 P1)。"""
from app.benchmarks import service
from app.services import note_service
note = asyncio.run(
note_service.create_note(
title="取消回归", markdown="向量数据库用于存储高维向量。", folder="", tags=["向量"]
)
)
cases = [
{
"case_id": f"c{i}",
"query": "向量数据库",
"expected_note_ids": [note.note_id],
"expected_block_ids": [note.blocks[0].block_id],
"citation_required": True,
}
for i in range(50)
]
_write_dataset("cancel-running-v1", cases)
async def _scenario():
run = await service.create_rag_run(
RAGRunRequest(dataset_id="cancel-running-v1", modes=[SearchMode.fts])
)
async def _cancel_after_start():
# 取消通过事件循环调度(独立 Task),而非同步直调,才能复现事件循环饥饿
while service.get_run(run.run_id).status == BenchmarkStatus.queued:
await asyncio.sleep(0)
service.cancel_run(run.run_id)
cancel_task = asyncio.create_task(_cancel_after_start())
finished = await service.wait_for_run(run.run_id)
await cancel_task
return finished
run = asyncio.run(_scenario())
assert run.status.value == "cancelled"
completed = sum(
1 for e in service.get_events(run.run_id) if e.event.value == "CaseCompleted"
)
assert completed < 50 # 未跑完全部样本,证明取消在样本边界生效
def test_block_only_annotation_resolves_note_and_scores() -> None:
"""仅标注 expected_block_ids 的样本应按块反查笔记评分,而非零分(审阅 P2)。"""
from app.services import note_service
note = asyncio.run(
note_service.create_note(
title="仅块标注", markdown="向量数据库存储高维向量。", folder="", tags=["向量"]
)
)
_write_dataset("block-only-v1", [{
"case_id": "c1",
"query": "向量数据库",
"expected_block_ids": [note.blocks[0].block_id],
"citation_required": False,
}])
run = _run(RAGRunRequest(dataset_id="block-only-v1", modes=[SearchMode.fts]))
assert run.status.value == "completed"
fts = run.metrics["fts"]
assert fts["hit_at_1"] == 1.0
assert fts["recall_at_k"] == 1.0
assert fts["mrr"] == 1.0
def test_sse_stream_ends_on_terminal_event_in_replay() -> None:
"""历史回放期间遇到终止事件时流应立即结束,而非进入实时队列永久等待(审阅 P2)。"""
from app import routes
from app.benchmarks import service
from app.contracts import BenchmarkEvent, BenchmarkEventType
run_id = "benchmark_sse_replay"
now = service._now()
# 模拟「回放期间运行完成」:run 仍为 runningsubscribe 返回非空队列),
# 但历史事件里已含 RunCompleted 终止事件。
service._runs[run_id] = BenchmarkRun(
run_id=run_id,
kind=BenchmarkKind.rag,
dataset_id="d",
dataset_hash="sha256:x",
status=BenchmarkStatus.running,
created_at=now,
)
service._events[run_id] = [
BenchmarkEvent(
event=BenchmarkEventType.run_started, run_id=run_id, sequence=0,
data={}, timestamp=now,
),
BenchmarkEvent(
event=BenchmarkEventType.run_completed, run_id=run_id, sequence=1,
data={}, timestamp=now,
),
]
try:
# 直调路由函数时 FastAPI 不解析 Query/Header 默认值,需显式传 None 覆盖 Header 哨兵
response = asyncio.run(
routes.benchmark_events(run_id, after_sequence=-1, last_event_id=None)
)
async def _collect() -> list[str]:
out: list[str] = []
async for chunk in response.body_iterator:
out.append(chunk)
return out
# 加超时防止回归(旧实现会永久挂起)
chunks = asyncio.run(asyncio.wait_for(_collect(), timeout=5))
finally:
service._forget(run_id)
events = [
line for chunk in chunks for line in chunk.splitlines() if line.startswith("event: ")
]
assert events == ["event: RunStarted", "event: RunCompleted"]