fix(provider): 修复索引事务回滚与工具名分片并同步主分支
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,198 @@
|
||||
"""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 BaseModel, Field, 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 = ""
|
||||
|
||||
|
||||
class _DatasetMeta(BaseModel):
|
||||
"""Dataset 元数据的最小校验模型。
|
||||
|
||||
list_datasets 用它逐文件校验元信息字段结构，把「合法 JSON 但字段类型错误」
|
||||
（如 cases: 42）这类损坏文件隔离掉，而不是让 len() 抛 TypeError 拖垮整个列表。
|
||||
"""
|
||||
|
||||
dataset_id: str = Field(min_length=1)
|
||||
kind: str = ""
|
||||
version: str = ""
|
||||
description: str = ""
|
||||
cases: list = Field(default_factory=list)
|
||||
|
||||
|
||||
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},
|
||||
)
|
||||
# citation_required=true 时必须声明 expected_block_ids，否则无法计算 Citation Hit Rate
|
||||
if parsed.citation_required and not parsed.expected_block_ids:
|
||||
raise ApiError(
|
||||
422,
|
||||
"BENCHMARK_DATASET_INVALID",
|
||||
f"Dataset case '{parsed.case_id}' requires expected_block_ids when citation_required is true.",
|
||||
{"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 内容）。
|
||||
|
||||
逐文件用 _DatasetMeta 校验元信息字段结构，单个损坏文件隔离跳过而非整体失败，
|
||||
保证列表接口健壮；损坏细节由 load_dataset 抛出。
|
||||
"""
|
||||
infos: list[BenchmarkDatasetInfo] = []
|
||||
for path in _dataset_files():
|
||||
try:
|
||||
raw, raw_bytes = _read_json(path)
|
||||
meta = _DatasetMeta.model_validate(raw)
|
||||
except (ApiError, ValidationError):
|
||||
continue
|
||||
if meta.kind not in ("", kind.value):
|
||||
continue
|
||||
infos.append(
|
||||
BenchmarkDatasetInfo(
|
||||
dataset_id=meta.dataset_id,
|
||||
kind=kind,
|
||||
version=meta.version,
|
||||
description=meta.description,
|
||||
case_count=len(meta.cases),
|
||||
content_hash=_content_hash(raw_bytes),
|
||||
)
|
||||
)
|
||||
return infos
|
||||
|
||||
|
||||
def load_dataset(dataset_id: str, kind: BenchmarkKind) -> RAGDataset:
|
||||
"""按文件名加载并校验数据集；找不到抛 BENCHMARK_DATASET_NOT_FOUND。
|
||||
|
||||
只读取与请求 dataset_id 同名的文件（{dataset_id}.json），无关文件的损坏（JSON 语法
|
||||
错误、UTF-8 解码错误、顶层非对象）不会阻断目标数据集加载；只有目标文件本身损坏
|
||||
才抛 BENCHMARK_DATASET_INVALID。按现有文件 stem 精确匹配，不拼接调用方传入的路径。
|
||||
"""
|
||||
for path in _dataset_files():
|
||||
if path.stem != dataset_id:
|
||||
continue
|
||||
raw, raw_bytes = _read_json(path)
|
||||
if not isinstance(raw, dict):
|
||||
raise ApiError(
|
||||
422,
|
||||
"BENCHMARK_DATASET_INVALID",
|
||||
"Dataset top-level must be a JSON object.",
|
||||
{"dataset_id": dataset_id, "path": path.name},
|
||||
)
|
||||
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,58 @@
|
||||
"""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。
|
||||
|
||||
结果先去重：检索结果是 Block 级，同一 Note 可能经多个 Block 重复出现，
|
||||
直接逐项计数会把同一 Note 算多次、导致 Recall 超过 1。
|
||||
"""
|
||||
if not expected:
|
||||
return 0.0
|
||||
return len(set(retrieved[:k]) & expected) / 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,158 @@
|
||||
"""RAG Benchmark Runner：调用检索引擎对数据集逐 Case 求值并聚合指标。
|
||||
|
||||
只读操作，直接复用 app.retrieval.engine 的 search()，不旁路检索链路。指标按
|
||||
(mode, case, repeat) 逐样本计算，再按 mode 聚合；失败样本按零分计入质量指标分母，
|
||||
避免把执行失败误判为检索质量（同时保留 total/successful/failed/failure_rate）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
|
||||
from app import repository
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BenchmarkCancelled(Exception):
|
||||
"""运行在 Case 之间被取消时抛出，用于中断后台执行并标记 cancelled。"""
|
||||
|
||||
|
||||
async def run_rag(
|
||||
dataset: RAGDataset,
|
||||
request: RAGRunRequest,
|
||||
on_case: Callable[[RAGCaseResult, int, int], None] | None = None,
|
||||
should_cancel: Callable[[], bool] | None = None,
|
||||
) -> tuple[dict[str, RAGMetrics], list[RAGCaseResult]]:
|
||||
"""执行 RAG Benchmark，返回 (按 mode 聚合的指标, 全部逐样本结果)。
|
||||
|
||||
on_case 在每个样本求值完成后回调 (result, done, total)，供上层更新进度与事件。
|
||||
should_cancel 在每个样本开始前被检查；返回 True 时抛出 BenchmarkCancelled 中断运行。
|
||||
"""
|
||||
total = len(request.modes) * len(dataset.cases) * request.repeat
|
||||
done = 0
|
||||
results: list[RAGCaseResult] = []
|
||||
|
||||
for mode in request.modes:
|
||||
for case in dataset.cases:
|
||||
expected_notes = _expected_notes(case)
|
||||
for repeat in range(request.repeat):
|
||||
# 让出事件循环：使运行中取消、SSE 进度与并发 API 请求能及时得到调度
|
||||
await asyncio.sleep(0)
|
||||
if should_cancel is not None and should_cancel():
|
||||
raise BenchmarkCancelled()
|
||||
result = await _evaluate_one(case, mode, request, repeat, expected_notes)
|
||||
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
|
||||
|
||||
|
||||
def _expected_notes(case: RAGDatasetCase) -> set[str]:
|
||||
"""返回笔记级期望 id；仅标注块 ID 时从块反查所属笔记，避免把标注缺失误判为检索失败。"""
|
||||
if case.expected_note_ids:
|
||||
return set(case.expected_note_ids)
|
||||
return {hit.note_id for hit in repository.get_block_hits(case.expected_block_ids)}
|
||||
|
||||
|
||||
async def _evaluate_one(
|
||||
case: RAGDatasetCase,
|
||||
mode: SearchMode,
|
||||
request: RAGRunRequest,
|
||||
repeat: int,
|
||||
expected_notes: set[str],
|
||||
) -> RAGCaseResult:
|
||||
search_request = SearchRequest(
|
||||
query=case.query,
|
||||
mode=mode,
|
||||
limit=request.retrieval.top_k,
|
||||
include_snippet=False,
|
||||
rrf_k=request.retrieval.rrf_k,
|
||||
rerank=request.retrieval.rerank,
|
||||
rerank_candidates=request.retrieval.rerank_candidates,
|
||||
score_threshold=request.retrieval.score_threshold,
|
||||
)
|
||||
start = time.perf_counter()
|
||||
try:
|
||||
response = await engine.search(search_request)
|
||||
latency_ms = (time.perf_counter() - start) * 1000.0
|
||||
except Exception as exc: # 单个样本失败不中断整个 Benchmark
|
||||
# 详细异常只进日志，公开响应只带项目错误码与安全消息，避免泄露路径/SQL 等敏感信息
|
||||
logger.warning(
|
||||
"RAG case evaluation failed: case=%s mode=%s", case.case_id, mode.value,
|
||||
exc_info=exc,
|
||||
)
|
||||
return RAGCaseResult(
|
||||
case_id=case.case_id,
|
||||
mode=mode,
|
||||
repeat=repeat,
|
||||
latency_ms=(time.perf_counter() - start) * 1000.0,
|
||||
citation_applicable=case.citation_required,
|
||||
error="RAG case evaluation failed.",
|
||||
error_code="BENCHMARK_CASE_EVALUATION_FAILED",
|
||||
)
|
||||
|
||||
retrieved_note_ids = [item.note_id for item in response.items]
|
||||
retrieved_block_ids = [item.block_id for item in response.items]
|
||||
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=case.citation_required,
|
||||
)
|
||||
|
||||
|
||||
def _aggregate(cases: list[RAGCaseResult], mode: SearchMode) -> RAGMetrics:
|
||||
samples = [c for c in cases if c.mode == mode]
|
||||
total = len(samples)
|
||||
failed = sum(1 for c in samples if c.error is not None)
|
||||
successful = total - failed
|
||||
if total == 0:
|
||||
return RAGMetrics()
|
||||
|
||||
# 延迟只统计成功样本；失败样本按零分计入质量指标分母，避免汇总虚高
|
||||
latencies = [c.latency_ms for c in samples if c.error is None]
|
||||
citation_samples = [c for c in samples if c.citation_applicable]
|
||||
return RAGMetrics(
|
||||
hit_at_1=m.mean([1.0 if (c.error is None and c.hit_at_1) else 0.0 for c in samples]),
|
||||
hit_at_5=m.mean([1.0 if (c.error is None and c.hit_at_5) else 0.0 for c in samples]),
|
||||
recall_at_k=m.mean([c.recall if c.error is None else 0.0 for c in samples]),
|
||||
mrr=m.mean([c.reciprocal_rank if c.error is None else 0.0 for c in samples]),
|
||||
citation_hit_rate=m.mean(
|
||||
[1.0 if (c.error is None and 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),
|
||||
total_cases=total,
|
||||
successful_cases=successful,
|
||||
failed_cases=failed,
|
||||
failure_rate=failed / total,
|
||||
)
|
||||
@@ -0,0 +1,348 @@
|
||||
"""Benchmark 服务：运行注册表、配置快照与报告组装。
|
||||
|
||||
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 logging
|
||||
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 BenchmarkCancelled, run_rag
|
||||
from app.config import get_settings
|
||||
from app.contracts import (
|
||||
BenchmarkEvent,
|
||||
BenchmarkEventType,
|
||||
BenchmarkKind,
|
||||
BenchmarkReport,
|
||||
BenchmarkRun,
|
||||
BenchmarkStatus,
|
||||
RAGCaseResult,
|
||||
RAGMetrics,
|
||||
RAGRunRequest,
|
||||
SearchMode,
|
||||
)
|
||||
from app.errors import ApiError
|
||||
from app.retrieval.engine import engine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_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
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _forget(run_id: str) -> None:
|
||||
"""移除一条 run 的全部内存态；仅在 run 处于终态时调用，避免打断活动任务。"""
|
||||
_runs.pop(run_id, None)
|
||||
_events.pop(run_id, None)
|
||||
_reports.pop(run_id, None)
|
||||
_tasks.pop(run_id, None)
|
||||
_subscribers.pop(run_id, None)
|
||||
_cancel_flags.pop(run_id, None)
|
||||
|
||||
|
||||
def _evict_terminal() -> bool:
|
||||
"""超过容量时淘汰最旧的终态 run；全部为活动 run 无法淘汰时返回 False。
|
||||
|
||||
绝不能删除仍在运行（queued/running）的 run：那会连带移除其 _cancel_flags 与
|
||||
_subscribers，使后台 Task 访问时抛出 KeyError。
|
||||
"""
|
||||
terminal = (BenchmarkStatus.completed, BenchmarkStatus.failed, BenchmarkStatus.cancelled)
|
||||
while len(_runs) >= MAX_RUNS:
|
||||
victim = next(
|
||||
(rid for rid, run in _runs.items() if run.status in terminal), None
|
||||
)
|
||||
if victim is None:
|
||||
return False
|
||||
_forget(victim)
|
||||
return True
|
||||
|
||||
|
||||
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,
|
||||
"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,
|
||||
}
|
||||
|
||||
|
||||
async def _validate_index_compatibility(request: RAGRunRequest) -> None:
|
||||
"""创建 RAG Run 前校验索引已建立且与当前 Embedding 模型/维度兼容。
|
||||
|
||||
空索引或不兼容索引会让所有模式得到全 0 指标，把环境/索引错误误判为检索质量差，
|
||||
故在创建时即拒绝，返回 BENCHMARK_INDEX_INCOMPATIBLE。
|
||||
"""
|
||||
stats = repository.stats()
|
||||
meta = repository.get_index_meta()
|
||||
needs_vector = any(m in (SearchMode.vector, SearchMode.hybrid) for m in request.modes)
|
||||
|
||||
reasons: list[str] = []
|
||||
if stats["blocks"] == 0:
|
||||
reasons.append("index is empty (no indexed blocks; run /api/index/rebuild first)")
|
||||
if needs_vector:
|
||||
if meta.get("embedding_model") != engine.embedding.model_id:
|
||||
reasons.append(
|
||||
f"embedding model mismatch: index={meta.get('embedding_model')!r}, "
|
||||
f"engine={engine.embedding.model_id!r}"
|
||||
)
|
||||
if meta.get("embedding_dim") != str(engine.embedding.dim):
|
||||
reasons.append(
|
||||
f"embedding dimension mismatch: index={meta.get('embedding_dim')!r}, "
|
||||
f"engine={engine.embedding.dim}"
|
||||
)
|
||||
if await engine.vector_store.count() == 0:
|
||||
reasons.append("vector index is empty")
|
||||
if reasons:
|
||||
raise ApiError(
|
||||
409,
|
||||
"BENCHMARK_INDEX_INCOMPATIBLE",
|
||||
"Benchmark index is not built or is incompatible with the current retrieval engine.",
|
||||
{"reasons": reasons},
|
||||
)
|
||||
|
||||
|
||||
async def create_rag_run(request: RAGRunRequest) -> BenchmarkRun:
|
||||
"""创建一次 RAG Benchmark，立即返回 queued 的 BenchmarkRun，由后台 Task 执行。"""
|
||||
dataset = datasets.load_dataset(request.dataset_id, BenchmarkKind.rag)
|
||||
await _validate_index_compatibility(request)
|
||||
|
||||
# 容量检查：先淘汰终态 run 腾空间；满容量且全为活动 run 时拒绝创建
|
||||
if not _evict_terminal():
|
||||
raise ApiError(
|
||||
429,
|
||||
"BENCHMARK_CAPACITY_EXCEEDED",
|
||||
"Benchmark run capacity exceeded; wait for active runs to finish.",
|
||||
{},
|
||||
)
|
||||
|
||||
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.queued,
|
||||
progress=0.0,
|
||||
config_snapshot=snapshot,
|
||||
created_at=_now(),
|
||||
)
|
||||
_runs[run_id] = 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])
|
||||
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]},
|
||||
)
|
||||
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,
|
||||
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(),
|
||||
}
|
||||
)
|
||||
emit(BenchmarkEventType.run_cancelled, {"status": BenchmarkStatus.cancelled.value})
|
||||
_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: # 单次运行失败不拖垮服务，记录错误后结束
|
||||
# 详细异常只进日志，公开响应仅带项目错误码与安全消息，避免泄露路径/SQL 等敏感信息
|
||||
logger.exception("Benchmark run failed: run_id=%s", run_id)
|
||||
_runs[run_id] = _runs[run_id].model_copy(
|
||||
update={
|
||||
"status": BenchmarkStatus.failed,
|
||||
"progress": 1.0,
|
||||
"error": "Benchmark run failed.",
|
||||
"error_code": "BENCHMARK_RUN_FAILED",
|
||||
"completed_at": _now(),
|
||||
}
|
||||
)
|
||||
emit(
|
||||
BenchmarkEventType.run_failed,
|
||||
{"error": "Benchmark run failed.", "error_code": "BENCHMARK_RUN_FAILED"},
|
||||
)
|
||||
_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="Benchmark run failed.",
|
||||
error_code="BENCHMARK_RUN_FAILED",
|
||||
)
|
||||
finish()
|
||||
return
|
||||
|
||||
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,
|
||||
)
|
||||
finish()
|
||||
|
||||
|
||||
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:
|
||||
"""取消运行：对 queued/running 设置取消标志，后台 Task 在 Case 边界检查后置为 cancelled。"""
|
||||
run = _runs.get(run_id)
|
||||
if run is None:
|
||||
return None
|
||||
if run.status in (BenchmarkStatus.queued, BenchmarkStatus.running):
|
||||
_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)
|
||||
@@ -24,6 +24,7 @@ class Settings:
|
||||
db_path: Path
|
||||
vault_path: Path
|
||||
attachments_path: Path
|
||||
benchmark_datasets_path: Path
|
||||
|
||||
|
||||
@lru_cache
|
||||
@@ -41,4 +42,7 @@ def get_settings() -> Settings:
|
||||
attachments_path=Path(
|
||||
os.getenv("APP_ATTACHMENTS_PATH", str(data_dir / "attachments"))
|
||||
),
|
||||
benchmark_datasets_path=Path(
|
||||
os.getenv("APP_BENCHMARK_DATASETS_PATH", str(data_dir / "benchmarks"))
|
||||
),
|
||||
)
|
||||
|
||||
@@ -144,6 +144,12 @@ class SearchRequest(Contract):
|
||||
limit: int = Field(default=20, ge=1, le=100)
|
||||
offset: int = Field(default=0, ge=0)
|
||||
include_snippet: bool = True
|
||||
# 检索调优参数（Benchmark 与 Skill 共用）：控制 RRF / 精排 / 候选池 / 分数阈值。
|
||||
# rerank_candidates=None 表示对全部候选精排（保留原有行为），Benchmark 传显式值。
|
||||
rrf_k: int = Field(default=60, ge=1)
|
||||
rerank: bool = True
|
||||
rerank_candidates: int | None = Field(default=None, ge=1)
|
||||
score_threshold: float = Field(default=0.0, ge=0.0)
|
||||
|
||||
|
||||
class Citation(Contract):
|
||||
@@ -1005,3 +1011,151 @@ class IndexJob(Contract):
|
||||
status: Literal["queued", "running", "completed", "failed"]
|
||||
scope: Literal["all", "notes", "vectors"]
|
||||
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，
|
||||
其余参数透传到 SearchRequest，由检索引擎实际执行。"""
|
||||
|
||||
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],
|
||||
min_length=1,
|
||||
)
|
||||
retrieval: RAGRetrievalConfig = Field(default_factory=RAGRetrievalConfig)
|
||||
repeat: int = Field(default=1, ge=1, le=10)
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
@field_validator("modes")
|
||||
@classmethod
|
||||
def _no_duplicate_modes(cls, value: list[SearchMode]) -> list[SearchMode]:
|
||||
if len(value) != len(set(value)):
|
||||
raise ValueError("modes must not contain duplicates")
|
||||
return value
|
||||
|
||||
|
||||
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
|
||||
# 样本构成：失败样本按零分计入质量指标，汇总不虚高；报告据此可知实际分母
|
||||
total_cases: int = 0
|
||||
successful_cases: int = 0
|
||||
failed_cases: int = 0
|
||||
failure_rate: 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
|
||||
error_code: 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"
|
||||
run_cancelled = "RunCancelled"
|
||||
|
||||
|
||||
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
|
||||
error_code: 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
|
||||
error_code: str | None = None
|
||||
|
||||
@@ -99,28 +99,25 @@ class OpenAICompatibleProvider(EventStreamingMixin, HTTPProviderMixin):
|
||||
raw = object_value(raw)
|
||||
index = token_count(raw.get("index", 0))
|
||||
function = object_value(raw.get("function") or {})
|
||||
call = calls.setdefault(index, {"id": "", "name": "", "arguments": "", "started": False})
|
||||
call = calls.setdefault(index, {"id": "", "name": "", "arguments": ""})
|
||||
if raw.get("id"):
|
||||
call["id"] = string_value(raw["id"])
|
||||
if function.get("name"):
|
||||
call["name"] += string_value(function["name"])
|
||||
fragment = string_value(function.get("arguments", ""))
|
||||
call["arguments"] += fragment
|
||||
if not call["started"] and call["name"]:
|
||||
call["id"] = call["id"] or f"call_{uuid4().hex}"
|
||||
call["started"] = True
|
||||
yield ModelEventType.tool_call_start, {"tool_call_id": call["id"], "name": call["name"]}
|
||||
fragment = call["arguments"]
|
||||
if call["started"] and fragment:
|
||||
yield ModelEventType.tool_call_delta, {"tool_call_id": call["id"], "arguments_delta": fragment}
|
||||
if choice.get("finish_reason"):
|
||||
finished = True
|
||||
if not finished:
|
||||
raise truncated_stream()
|
||||
for call in calls.values():
|
||||
if not call["started"]:
|
||||
if not call["name"]:
|
||||
raise invalid_response()
|
||||
decode_tool_arguments(call["arguments"] or "{}")
|
||||
# A name can span multiple chunks; publish only the complete identity.
|
||||
call["id"] = call["id"] or f"call_{uuid4().hex}"
|
||||
yield ModelEventType.tool_call_start, {"tool_call_id": call["id"], "name": call["name"]}
|
||||
yield ModelEventType.tool_call_delta, {"tool_call_id": call["id"], "arguments_delta": call["arguments"] or "{}"}
|
||||
yield ModelEventType.tool_call_end, {"tool_call_id": call["id"]}
|
||||
|
||||
async def list_models(self) -> list[ModelInfo]:
|
||||
|
||||
+88
-18
@@ -273,11 +273,15 @@ def update_note_location(
|
||||
raise LookupError(note_id)
|
||||
|
||||
|
||||
def fts_search_page(
|
||||
*,
|
||||
_FTS_FROM = """
|
||||
FROM blocks_fts
|
||||
JOIN blocks AS b ON b.block_id = blocks_fts.block_id
|
||||
JOIN notes AS n ON n.note_id = b.note_id
|
||||
"""
|
||||
|
||||
|
||||
def _fts_where(
|
||||
match: str,
|
||||
limit: int,
|
||||
offset: int,
|
||||
folders: list[str],
|
||||
note_ids: list[str],
|
||||
tags: list[str],
|
||||
@@ -285,8 +289,11 @@ def fts_search_page(
|
||||
created_to: datetime | None,
|
||||
updated_from: datetime | None,
|
||||
updated_to: datetime | None,
|
||||
) -> tuple[list[FtsHit], int]:
|
||||
"""执行带元数据过滤的 FTS 精确分页，并返回过滤后的完整命中数。"""
|
||||
) -> tuple[str, list[object]]:
|
||||
"""构建 FTS 过滤 WHERE 子句（不含 WHERE 关键字），返回 (where_sql, params)。
|
||||
|
||||
fts_search_page 与 fts_score_bounds 共用，保证计数与取数口径一致。
|
||||
"""
|
||||
where = ["blocks_fts MATCH ?"]
|
||||
params: list[object] = [match]
|
||||
|
||||
@@ -317,22 +324,44 @@ def fts_search_page(
|
||||
where.append(f"julianday({column}) <= julianday(?)")
|
||||
params.append(_iso(upper))
|
||||
|
||||
from_sql = """
|
||||
FROM blocks_fts
|
||||
JOIN blocks AS b ON b.block_id = blocks_fts.block_id
|
||||
JOIN notes AS n ON n.note_id = b.note_id
|
||||
return " AND ".join(where), params
|
||||
|
||||
|
||||
def fts_search_page(
|
||||
*,
|
||||
match: str,
|
||||
limit: int,
|
||||
offset: int,
|
||||
folders: list[str],
|
||||
note_ids: list[str],
|
||||
tags: list[str],
|
||||
created_from: datetime | None,
|
||||
created_to: datetime | None,
|
||||
updated_from: datetime | None,
|
||||
updated_to: datetime | None,
|
||||
bm25_max: float | None = None,
|
||||
) -> tuple[list[FtsHit], int]:
|
||||
"""执行带元数据过滤的 FTS 精确分页，并返回过滤后的完整命中数。
|
||||
|
||||
bm25_max 非空时按 bm25 截止值过滤（用于阈值过滤的精确分页），计数与取数同口径。
|
||||
"""
|
||||
where_sql = " AND ".join(where)
|
||||
where_sql, params = _fts_where(
|
||||
match, folders, note_ids, tags,
|
||||
created_from, created_to, updated_from, updated_to,
|
||||
)
|
||||
if bm25_max is not None:
|
||||
where_sql += " AND bm25(blocks_fts) <= ?"
|
||||
params.append(bm25_max)
|
||||
|
||||
conn = connect()
|
||||
try:
|
||||
total = conn.execute(
|
||||
f"SELECT COUNT(*) {from_sql} WHERE {where_sql}", params
|
||||
f"SELECT COUNT(*) {_FTS_FROM} WHERE {where_sql}", params
|
||||
).fetchone()[0]
|
||||
rows = conn.execute(
|
||||
f"""
|
||||
SELECT blocks_fts.block_id, blocks_fts.note_id, bm25(blocks_fts) AS rank
|
||||
{from_sql}
|
||||
{_FTS_FROM}
|
||||
WHERE {where_sql}
|
||||
ORDER BY rank
|
||||
LIMIT ? OFFSET ?
|
||||
@@ -348,6 +377,45 @@ def fts_search_page(
|
||||
conn.close()
|
||||
|
||||
|
||||
def fts_score_bounds(
|
||||
*,
|
||||
match: str,
|
||||
folders: list[str],
|
||||
note_ids: list[str],
|
||||
tags: list[str],
|
||||
created_from: datetime | None,
|
||||
created_to: datetime | None,
|
||||
updated_from: datetime | None,
|
||||
updated_to: datetime | None,
|
||||
) -> tuple[float, float] | None:
|
||||
"""返回 metadata 过滤后的 FTS 命中集里 bm25 的 (min, max)，无命中时返回 None。
|
||||
|
||||
用于阈值过滤：min-max 归一化是 bm25 的线性函数，据此可把阈值换算为 bm25 截止值。
|
||||
"""
|
||||
where_sql, params = _fts_where(
|
||||
match, folders, note_ids, tags,
|
||||
created_from, created_to, updated_from, updated_to,
|
||||
)
|
||||
conn = connect()
|
||||
try:
|
||||
# bm25() 不能作为聚合函数参数，也不能用在被聚合的子查询里；改用 ORDER BY 取首尾两行
|
||||
lo_row = conn.execute(
|
||||
f"SELECT bm25(blocks_fts) AS rank {_FTS_FROM} WHERE {where_sql}"
|
||||
" ORDER BY rank ASC LIMIT 1",
|
||||
params,
|
||||
).fetchone()
|
||||
if lo_row is None or lo_row["rank"] is None:
|
||||
return None
|
||||
hi_row = conn.execute(
|
||||
f"SELECT bm25(blocks_fts) AS rank {_FTS_FROM} WHERE {where_sql}"
|
||||
" ORDER BY rank DESC LIMIT 1",
|
||||
params,
|
||||
).fetchone()
|
||||
return (float(lo_row["rank"]), float(hi_row["rank"]))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_block_hits(block_ids: list[str]) -> list[BlockHit]:
|
||||
if not block_ids:
|
||||
return []
|
||||
@@ -392,16 +460,18 @@ def get_index_meta() -> dict[str, str]:
|
||||
conn.close()
|
||||
|
||||
|
||||
def clear_all() -> None:
|
||||
"""清空元数据、Block 与 FTS5（重建索引用，向量由 VectorStore.clear 处理）。"""
|
||||
conn = connect()
|
||||
def clear_all(*, conn: sqlite3.Connection | None = None) -> None:
|
||||
"""Clear rebuildable metadata using the caller's transaction when provided."""
|
||||
owns = conn is None
|
||||
conn = conn or connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
with transaction(conn) if owns else nullcontext():
|
||||
conn.execute("DELETE FROM blocks_fts")
|
||||
conn.execute("DELETE FROM blocks")
|
||||
conn.execute("DELETE FROM notes")
|
||||
finally:
|
||||
conn.close()
|
||||
if owns:
|
||||
conn.close()
|
||||
|
||||
|
||||
def stats() -> dict[str, int]:
|
||||
|
||||
@@ -19,6 +19,7 @@ class EmbeddingProvider(Protocol):
|
||||
"""统一 Embedding 接口（与文档一致）。"""
|
||||
|
||||
model_id: str
|
||||
version: str
|
||||
dim: int
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
@@ -33,6 +34,7 @@ class HashEmbeddingProvider:
|
||||
"""
|
||||
|
||||
model_id = "hash-v1"
|
||||
version = "1"
|
||||
dim = EMBEDDING_DIM
|
||||
|
||||
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
|
||||
@@ -62,7 +62,7 @@ class RetrievalEngine:
|
||||
# 候选池至少覆盖本次请求的 offset+limit，保证分页能取到目标页；设上限防内存失控
|
||||
window = min(request.offset + request.limit, MAX_CANDIDATE_POOL)
|
||||
pool_size = max(CANDIDATE_POOL, window)
|
||||
# 带过滤时放大召回；FTS 则一次性取全量命中（≤FTS_FETCH_LIMIT）避免截断漏召回
|
||||
# 带过滤时放大召回，缓解「先截断候选池再过滤」造成的漏召回
|
||||
recall = min(pool_size * OVERSCAN_FACTOR, MAX_CANDIDATE_POOL) if has_filters else pool_size
|
||||
|
||||
# 1. 按模式收集候选（FTS 与 Vector 各产出「按相关性降序」的 block_id 列表）
|
||||
@@ -98,7 +98,7 @@ class RetrievalEngine:
|
||||
elif request.mode == SearchMode.vector:
|
||||
candidate_scores = vec_scores
|
||||
else: # hybrid：RRF 融合
|
||||
candidate_scores = rrf_fuse([fts_ranked, vec_ranked])
|
||||
candidate_scores = rrf_fuse([fts_ranked, vec_ranked], k=request.rrf_k)
|
||||
|
||||
if not candidate_scores:
|
||||
return self._empty(request)
|
||||
@@ -111,14 +111,23 @@ class RetrievalEngine:
|
||||
if not filtered:
|
||||
return self._empty(request)
|
||||
|
||||
# 4. 排序 / 精排
|
||||
# 4. 排序 / 精排：hybrid 先按融合分预排序，再对前 rerank_candidates 个候选做精排，
|
||||
# 剩余候选按融合分排在精排结果之后；rerank=False 时跳过精排直接按融合分排序。
|
||||
if request.mode == SearchMode.hybrid:
|
||||
candidates = [
|
||||
RankedCandidate(block_id=h.block_id, score=candidate_scores[h.block_id], text=h.content)
|
||||
for h in filtered
|
||||
]
|
||||
ranked = await self.reranker.rerank(request.query, candidates)
|
||||
ordered = [(c.block_id, c.score) for c in ranked]
|
||||
pre_sorted = sorted(filtered, key=lambda h: -candidate_scores[h.block_id])
|
||||
if request.rerank:
|
||||
limit = request.rerank_candidates
|
||||
pool = pre_sorted if limit is None else pre_sorted[:limit]
|
||||
rest = [] if limit is None else pre_sorted[limit:]
|
||||
candidates = [
|
||||
RankedCandidate(block_id=h.block_id, score=candidate_scores[h.block_id], text=h.content)
|
||||
for h in pool
|
||||
]
|
||||
ranked = await self.reranker.rerank(request.query, candidates)
|
||||
ordered = [(c.block_id, c.score) for c in ranked]
|
||||
ordered += [(h.block_id, candidate_scores[h.block_id]) for h in rest]
|
||||
else:
|
||||
ordered = [(h.block_id, candidate_scores[h.block_id]) for h in pre_sorted]
|
||||
else:
|
||||
ordered = sorted(
|
||||
((h.block_id, candidate_scores[h.block_id]) for h in filtered),
|
||||
@@ -126,8 +135,10 @@ class RetrievalEngine:
|
||||
)
|
||||
|
||||
ordered = normalize_scores(ordered)
|
||||
# score_threshold：归一化后过滤低分结果（默认 0 不过滤）
|
||||
ordered = [(bid, score) for bid, score in ordered if score >= request.score_threshold]
|
||||
|
||||
# 5. 分页：total = 过滤后候选集大小。fts 已取全量（≤FTS_FETCH_LIMIT）故为真实命中数；
|
||||
# 5. 分页：total = 过滤后候选集大小。fts 走数据库精确分页，total 为真实命中数；
|
||||
# vector/hybrid 为 KNN 候选集，无全局 total。
|
||||
total = len(ordered)
|
||||
page = ordered[request.offset : request.offset + request.limit]
|
||||
@@ -140,11 +151,41 @@ class RetrievalEngine:
|
||||
)
|
||||
|
||||
def _search_fts(self, request: SearchRequest) -> SearchResponse:
|
||||
"""FTS 专用路径：过滤、COUNT 与分页全部在 SQLite 中完成。"""
|
||||
"""FTS 专用路径：在数据库侧完成过滤、计数与分页，不取全量后再截断。
|
||||
|
||||
阈值过滤时，min-max 归一化是 bm25 的线性函数，据此把 score_threshold 换算为
|
||||
bm25 截止值（bm25_max），使过滤、计数与分页口径一致；无阈值时走数据库原生分页，
|
||||
total 始终为过滤后的真实命中数，不再受固定截断影响。
|
||||
"""
|
||||
match = match_query(request.query)
|
||||
if not match:
|
||||
return self._empty(request)
|
||||
|
||||
bounds = repository.fts_score_bounds(
|
||||
match=match,
|
||||
folders=request.folders,
|
||||
note_ids=request.note_ids,
|
||||
tags=request.tags,
|
||||
created_from=request.created_from,
|
||||
created_to=request.created_to,
|
||||
updated_from=request.updated_from,
|
||||
updated_to=request.updated_to,
|
||||
)
|
||||
if bounds is None:
|
||||
return self._empty(request)
|
||||
|
||||
lo, hi = bounds
|
||||
span = hi - lo
|
||||
bm25_max: float | None = None
|
||||
if request.score_threshold > 0:
|
||||
if span == 0:
|
||||
# 全部命中 bm25 相同，归一化后皆为 1.0；阈值超过 1.0 时无命中
|
||||
if request.score_threshold > 1.0:
|
||||
return self._empty(request)
|
||||
else:
|
||||
# norm = (hi - bm25) / span；norm >= threshold ⟺ bm25 <= hi - threshold * span
|
||||
bm25_max = hi - request.score_threshold * span
|
||||
|
||||
fts_hits, total = repository.fts_search_page(
|
||||
match=match,
|
||||
limit=request.limit,
|
||||
@@ -156,19 +197,29 @@ class RetrievalEngine:
|
||||
created_to=request.created_to,
|
||||
updated_from=request.updated_from,
|
||||
updated_to=request.updated_to,
|
||||
bm25_max=bm25_max,
|
||||
)
|
||||
if not fts_hits:
|
||||
# 本页无结果：offset 越过末页时 total 仍为真实命中数（>0），需保留而非归零
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
mode=request.mode,
|
||||
items=[],
|
||||
page=PageMeta(total=total, limit=request.limit, offset=request.offset),
|
||||
)
|
||||
|
||||
hits = {h.block_id: h for h in repository.get_block_hits([hit.block_id for hit in fts_hits])}
|
||||
ordered = normalize_scores(
|
||||
[(hit.block_id, -hit.bm25) for hit in fts_hits if hit.block_id in hits]
|
||||
)
|
||||
items = [self._build_result(hits[block_id], request, score) for block_id, score in ordered]
|
||||
# 分数按全局 bm25 上下界归一化（与取全量后 normalize_scores 等价），保证跨页一致
|
||||
span = hi - lo
|
||||
if span == 0:
|
||||
ordered = [(hit.block_id, 1.0) for hit in fts_hits]
|
||||
else:
|
||||
ordered = [(hit.block_id, round((hi - hit.bm25) / span, 6)) for hit in fts_hits]
|
||||
hits = {h.block_id: h for h in repository.get_block_hits([bid for bid, _ in ordered])}
|
||||
items = [
|
||||
self._build_result(hits[block_id], request, score)
|
||||
for block_id, score in ordered
|
||||
if block_id in hits
|
||||
]
|
||||
return SearchResponse(
|
||||
query=request.query,
|
||||
mode=request.mode,
|
||||
|
||||
@@ -24,6 +24,7 @@ class RerankerProvider(Protocol):
|
||||
"""统一 Reranker 接口：输入候选块，输出按相关性重排后的候选块。"""
|
||||
|
||||
model_id: str
|
||||
version: str
|
||||
|
||||
async def rerank(self, query: str, candidates: list[RankedCandidate]) -> list[RankedCandidate]: ...
|
||||
|
||||
@@ -32,6 +33,7 @@ class LexicalReranker:
|
||||
"""轻量精排：query 与块正文的词面重叠度，与归一化后的原始分数加权求和。"""
|
||||
|
||||
model_id = "lexical-v1"
|
||||
version = "1"
|
||||
|
||||
def __init__(self, lexical_weight: float = 0.5) -> None:
|
||||
self.lexical_weight = lexical_weight
|
||||
|
||||
@@ -35,6 +35,7 @@ class VectorStore(Protocol):
|
||||
async def upsert(self, records: list[VectorRecord]) -> None: ...
|
||||
async def delete(self, ids: list[str]) -> None: ...
|
||||
async def search(self, vector: list[float], *, top_k: int) -> list[VectorHit]: ...
|
||||
async def count(self) -> int: ...
|
||||
|
||||
|
||||
class SqliteVecStore:
|
||||
@@ -85,10 +86,19 @@ class SqliteVecStore:
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
async def clear(self) -> None:
|
||||
conn = connect()
|
||||
async def clear(self, *, conn: sqlite3.Connection | None = None) -> None:
|
||||
owns = conn is None
|
||||
conn = conn or connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
with transaction(conn) if owns else nullcontext():
|
||||
conn.execute("DELETE FROM vec_blocks")
|
||||
finally:
|
||||
if owns:
|
||||
conn.close()
|
||||
|
||||
async def count(self) -> int:
|
||||
conn = connect()
|
||||
try:
|
||||
return conn.execute("SELECT COUNT(*) FROM vec_blocks").fetchone()[0]
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
@@ -15,6 +15,14 @@ from app.contracts import (
|
||||
AgentRunListResponse,
|
||||
AgentTraceResponse,
|
||||
ChatRequest,
|
||||
BenchmarkDatasetListResponse,
|
||||
BenchmarkEventType,
|
||||
BenchmarkKind,
|
||||
BenchmarkReport,
|
||||
BenchmarkRun,
|
||||
BenchmarkRunListResponse,
|
||||
BenchmarkStatus,
|
||||
RAGRunRequest,
|
||||
CredentialStatus,
|
||||
CredentialWriteRequest,
|
||||
ExtensionInstallRequest,
|
||||
@@ -85,6 +93,10 @@ from app.contracts import (
|
||||
WorkspaceOpenRequest,
|
||||
WorkspaceSnapshot,
|
||||
)
|
||||
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.errors import ApiError
|
||||
from app.extensions import ExtensionError
|
||||
from app.extensions.mcp_registry import McpRegistryError
|
||||
@@ -1133,3 +1145,168 @@ async def get_index_job(job_id: str) -> IndexJob:
|
||||
404, "RESOURCE_NOT_FOUND", "index job not found", {"job_id": job_id}
|
||||
)
|
||||
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.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="accepted",
|
||||
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),
|
||||
last_event_id: str | None = Header(default=None, alias="Last-Event-ID"),
|
||||
) -> 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}
|
||||
)
|
||||
|
||||
# SSE 断线重连：Last-Event-ID 优先于 after_sequence，用于从上次收到的事件继续
|
||||
cursor = after_sequence
|
||||
if last_event_id is not None:
|
||||
try:
|
||||
cursor = int(last_event_id)
|
||||
except ValueError as exc:
|
||||
raise ApiError(
|
||||
400,
|
||||
"BENCHMARK_EVENT_CURSOR_INVALID",
|
||||
"Last-Event-ID must be an integer sequence.",
|
||||
{"last_event_id": last_event_id},
|
||||
) from exc
|
||||
if cursor < -1:
|
||||
raise ApiError(
|
||||
400,
|
||||
"BENCHMARK_EVENT_CURSOR_INVALID",
|
||||
"Last-Event-ID must be greater than or equal to -1.",
|
||||
)
|
||||
|
||||
async def stream() -> AsyncIterator[str]:
|
||||
# 先订阅（保证订阅之后产生的事件也能收到），再回放历史事件，最后实时输出新事件
|
||||
terminal = (
|
||||
BenchmarkEventType.run_completed,
|
||||
BenchmarkEventType.run_failed,
|
||||
BenchmarkEventType.run_cancelled,
|
||||
)
|
||||
queue = benchmark_service.subscribe(run_id)
|
||||
try:
|
||||
last_sequence = cursor
|
||||
# 回放按订阅时刻的快照长度遍历，避免列表在回放期间被追加；终止事件同样要结束流，
|
||||
# 防止回放完成后进入实时队列却因序号去重跳过同一终止事件而永久等待。
|
||||
history = benchmark_service.get_events(run_id)
|
||||
for index in range(len(history)):
|
||||
event = history[index]
|
||||
if event.sequence <= cursor:
|
||||
continue
|
||||
yield as_sse(event.event.value, event.model_dump_json(), event_id=event.sequence)
|
||||
last_sequence = event.sequence
|
||||
if event.event in terminal:
|
||||
return
|
||||
if queue is None:
|
||||
return
|
||||
while True:
|
||||
event = await queue.get()
|
||||
if event.sequence <= last_sequence:
|
||||
continue
|
||||
yield as_sse(event.event.value, event.model_dump_json(), event_id=event.sequence)
|
||||
last_sequence = event.sequence
|
||||
if event.event in terminal:
|
||||
return
|
||||
finally:
|
||||
if queue is not None:
|
||||
benchmark_service.unsubscribe(run_id, queue)
|
||||
|
||||
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
|
||||
|
||||
@@ -6,7 +6,6 @@ MVP 阶段重建是同步的（数据量小），完成后直接返回 completed
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
@@ -16,8 +15,8 @@ from app.config import get_settings
|
||||
from app.contracts import IndexJob, IndexRebuildRequest, IndexStatus
|
||||
from app.errors import ApiError
|
||||
from app.knowledge.parser import parse_note
|
||||
from app.services.note_service import index_note
|
||||
from app.services import task_service
|
||||
from app.services.note_service import index_note, prepare_note_index
|
||||
from app.database.db import connect, transaction
|
||||
from app.services.coordination import serialized_vault_mutation
|
||||
from app.retrieval.vectorstore import SqliteVecStore
|
||||
|
||||
@@ -74,18 +73,7 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
|
||||
{"scope": request.scope, "note_ids": request.note_ids},
|
||||
)
|
||||
|
||||
# 先扫描到内存（失败不会清旧索引），再快照旧库用于失败回滚
|
||||
docs = _scan_vault()
|
||||
settings = get_settings()
|
||||
database_existed = settings.db_path.exists()
|
||||
task_note_links = task_service.note_links() if database_existed else {}
|
||||
backup_path = (
|
||||
settings.db_path.with_name(f"{settings.db_path.name}.{job_id}.bak")
|
||||
if database_existed
|
||||
else None
|
||||
)
|
||||
if backup_path is not None:
|
||||
shutil.copy2(settings.db_path, backup_path)
|
||||
|
||||
_active_job_id = job_id
|
||||
_last_error = None
|
||||
@@ -94,23 +82,34 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
|
||||
created_at=datetime.now(timezone.utc),
|
||||
))
|
||||
try:
|
||||
# Deleting blocks also cascades every space in routed_block_vectors;
|
||||
# index_note repopulates only the currently successful API space.
|
||||
repository.clear_all()
|
||||
await vector_store.clear()
|
||||
prepared_notes = []
|
||||
for rel, folder, markdown, created, updated in docs:
|
||||
parsed = parse_note(
|
||||
markdown=markdown, file_path=rel, folder=folder, tags=None,
|
||||
created_at=created, updated_at=updated,
|
||||
)
|
||||
await index_note(parsed)
|
||||
task_service.restore_note_links(task_note_links)
|
||||
prepared_notes.append((parsed, await prepare_note_index(parsed)))
|
||||
# All network/model awaits precede the transaction. The concrete SQLite
|
||||
# methods below complete synchronously despite their async interfaces.
|
||||
conn = connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
task_note_links = dict(conn.execute(
|
||||
"SELECT task_id, note_id FROM tasks WHERE note_id IS NOT NULL"
|
||||
).fetchall())
|
||||
repository.clear_all(conn=conn)
|
||||
await vector_store.clear(conn=conn)
|
||||
for parsed, prepared in prepared_notes:
|
||||
await index_note(parsed, prepared=prepared, conn=conn)
|
||||
for task_id, note_id in task_note_links.items():
|
||||
conn.execute(
|
||||
"UPDATE tasks SET note_id = ? WHERE task_id = ? "
|
||||
"AND EXISTS (SELECT 1 FROM notes WHERE note_id = ?)",
|
||||
(note_id, task_id, note_id),
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
except BaseException as exc:
|
||||
# 重建失败：恢复旧索引，避免留下半成品；记录 failed 任务后向上抛
|
||||
if backup_path is not None and backup_path.exists():
|
||||
shutil.copy2(backup_path, settings.db_path)
|
||||
elif not database_existed:
|
||||
settings.db_path.unlink(missing_ok=True)
|
||||
_remember_job(IndexJob(
|
||||
job_id=job_id, status="failed", scope=request.scope,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
@@ -119,8 +118,6 @@ async def rebuild(request: IndexRebuildRequest) -> IndexJob:
|
||||
raise
|
||||
finally:
|
||||
_active_job_id = None
|
||||
if backup_path is not None:
|
||||
backup_path.unlink(missing_ok=True)
|
||||
|
||||
job = IndexJob(job_id=job_id, status="completed", scope=request.scope, created_at=datetime.now(timezone.utc))
|
||||
_remember_job(job)
|
||||
|
||||
@@ -6,6 +6,8 @@ Markdown 文件是笔记正文的持久化载体（Vault），SQLite/FTS5/向量
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
from contextlib import nullcontext
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
@@ -72,21 +74,34 @@ def _delete_markdown(rel_path: str) -> None:
|
||||
path.unlink()
|
||||
|
||||
|
||||
async def index_note(parsed: ParsedNote) -> None:
|
||||
PreparedIndex = tuple[list[list[float]], routed_vectors.RemoteEmbeddings | None]
|
||||
|
||||
|
||||
async def prepare_note_index(parsed: ParsedNote) -> PreparedIndex:
|
||||
"""Compute vectors before opening a write transaction (including API I/O)."""
|
||||
texts = [block.content for block in parsed.blocks]
|
||||
vectors = await embedding.embed_documents(texts)
|
||||
remote = await routed_vectors.embed_remote(texts)
|
||||
return vectors, remote
|
||||
|
||||
|
||||
async def index_note(
|
||||
parsed: ParsedNote, *, prepared: PreparedIndex | None = None,
|
||||
conn: sqlite3.Connection | None = None,
|
||||
) -> None:
|
||||
"""把解析结果写入元数据 + FTS5 + 向量（三层可重建索引），单事务保证原子性。
|
||||
|
||||
元数据与向量在同一连接、同一事务内提交，避免「新元数据已提交、向量写入失败」的
|
||||
半提交状态。替换元数据时拿到旧 block_id：清理已删除/内容变化的旧向量，只为新增
|
||||
block 写向量（内容未变的 block 其向量仍有效，无需重复写入）。
|
||||
"""
|
||||
texts = [block.content for block in parsed.blocks]
|
||||
vectors = await embedding.embed_documents(texts)
|
||||
# Network I/O stays outside the write transaction. The hash index remains
|
||||
# complete even when the optional API route fails or changes vector spaces.
|
||||
remote = await routed_vectors.embed_remote(texts)
|
||||
conn = connect()
|
||||
if conn is not None and prepared is None:
|
||||
raise ValueError("Prepare embeddings before supplying a write connection")
|
||||
vectors, remote = prepared if prepared is not None else await prepare_note_index(parsed)
|
||||
owns = conn is None
|
||||
conn = conn or connect()
|
||||
try:
|
||||
with transaction(conn):
|
||||
with transaction(conn) if owns else nullcontext():
|
||||
old_block_ids = repository.replace_note_metadata(
|
||||
conn=conn,
|
||||
note_id=parsed.note_id,
|
||||
@@ -116,7 +131,8 @@ async def index_note(parsed: ParsedNote) -> None:
|
||||
conn=conn,
|
||||
)
|
||||
finally:
|
||||
conn.close()
|
||||
if owns:
|
||||
conn.close()
|
||||
|
||||
|
||||
@serialized_vault_mutation
|
||||
|
||||
@@ -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,566 @@
|
||||
"""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"]["version"]
|
||||
assert snapshot["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_applicable（citation_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 仍为 running（subscribe 返回非空队列），
|
||||
# 但历史事件里已含 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"]
|
||||
@@ -82,6 +82,31 @@ async def collect(iterator):
|
||||
return [event async for event in iterator]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["lookup", "notes.search"])
|
||||
def test_compatible_split_tool_name_preserves_identity(name):
|
||||
from app.providers.tool_names import prepare_tool_names
|
||||
req = request()
|
||||
req.tools[0].name = name
|
||||
wire, _ = prepare_tool_names(req)
|
||||
alias = wire.tools[0].name
|
||||
|
||||
def handler(_):
|
||||
return httpx.Response(200, content=sse(
|
||||
{"choices": [{"delta": {"tool_calls": [{"index": 0, "id": "call_1",
|
||||
"function": {"name": alias[:3], "arguments": ""}}]}}]},
|
||||
{"choices": [{"delta": {"tool_calls": [{"index": 0,
|
||||
"function": {"name": alias[3:], "arguments": '{"query":"x"}'}}]},
|
||||
"finish_reason": "tool_calls"}]},
|
||||
{"type": "[DONE]"},
|
||||
))
|
||||
|
||||
events = asyncio.run(collect(provider("compatible", handler).stream(req)))
|
||||
assert [e.data["name"] for e in events if e.event == E.tool_call_start] == [name]
|
||||
assert json.loads("".join(e.data["arguments_delta"] for e in events
|
||||
if e.event == E.tool_call_delta)) == {"query": "x"}
|
||||
assert events[-1].data["status"] == "completed"
|
||||
|
||||
|
||||
def sse(*events):
|
||||
return "".join(
|
||||
f"event: {event.get('type', 'message')}\r\ndata: {json.dumps(event, ensure_ascii=False)}\r\n\r\n"
|
||||
|
||||
@@ -436,6 +436,83 @@ def test_fts_pagination_is_not_truncated_at_one_thousand(vault) -> None:
|
||||
assert len(response.items) == 10
|
||||
|
||||
|
||||
def test_fts_score_threshold_filters_before_total(vault) -> None:
|
||||
"""score_threshold 先于计数与分页生效：total 反映过滤后数量，与 items 一致。
|
||||
|
||||
高阈值过滤掉全部结果时 total==0 且 items 为空，杜绝「空页但 total>0」的
|
||||
不一致（审阅 P2-7）。
|
||||
"""
|
||||
from app.retrieval.engine import engine
|
||||
from app.services import note_service
|
||||
|
||||
# 10 个 block，含「目标」次数递增，bm25 分数各异，min-max 归一化后分数落在 [0,1]
|
||||
markdown = "\n\n".join(f"{'目标' * i} 分隔内容" for i in range(1, 11))
|
||||
asyncio.run(
|
||||
note_service.create_note(title="阈值过滤", markdown=markdown, folder="", tags=[])
|
||||
)
|
||||
|
||||
all_hits = asyncio.run(
|
||||
engine.search(
|
||||
SearchRequest(query="目标", mode=SearchMode.fts, limit=20, score_threshold=0.0)
|
||||
)
|
||||
)
|
||||
filtered = asyncio.run(
|
||||
engine.search(
|
||||
SearchRequest(query="目标", mode=SearchMode.fts, limit=20, score_threshold=0.5)
|
||||
)
|
||||
)
|
||||
none = asyncio.run(
|
||||
engine.search(
|
||||
SearchRequest(query="目标", mode=SearchMode.fts, limit=20, score_threshold=2.0)
|
||||
)
|
||||
)
|
||||
|
||||
assert all_hits.page.total >= 10
|
||||
assert 0 < filtered.page.total < all_hits.page.total # 阈值过滤掉部分而非全部
|
||||
assert filtered.page.total == len(filtered.items)
|
||||
assert none.page.total == 0
|
||||
assert none.items == []
|
||||
|
||||
|
||||
def test_fts_offset_beyond_end_reports_real_total(vault) -> None:
|
||||
"""offset 越过末页时 items 为空，但 total 仍为真实命中数而非归零。"""
|
||||
from app.retrieval.engine import engine
|
||||
from app.services import note_service
|
||||
|
||||
asyncio.run(
|
||||
note_service.create_note(title="越界分页", markdown="检索 检索 检索 检索", folder="", tags=[])
|
||||
)
|
||||
|
||||
resp = asyncio.run(
|
||||
engine.search(SearchRequest(query="检索", mode=SearchMode.fts, limit=10, offset=100))
|
||||
)
|
||||
assert resp.page.total >= 1
|
||||
assert resp.items == []
|
||||
|
||||
|
||||
def test_fts_not_truncated_at_five_thousand(vault) -> None:
|
||||
"""FTS 结果不再被 5000 条上限截断：>5000 命中时 total 为真实计数，末页仍可访问。"""
|
||||
from app.retrieval.engine import engine
|
||||
from app.services import note_service
|
||||
|
||||
markdown = "\n\n".join(f"共同词 q{i}" for i in range(5010))
|
||||
asyncio.run(
|
||||
note_service.create_note(title="五千条分页", markdown=markdown, folder="", tags=[])
|
||||
)
|
||||
|
||||
first = asyncio.run(
|
||||
engine.search(SearchRequest(query="共同词", mode=SearchMode.fts, limit=10, offset=0))
|
||||
)
|
||||
assert first.page.total == 5010
|
||||
assert len(first.items) == 10
|
||||
|
||||
last = asyncio.run(
|
||||
engine.search(SearchRequest(query="共同词", mode=SearchMode.fts, limit=10, offset=5005))
|
||||
)
|
||||
assert last.page.total == 5010
|
||||
assert len(last.items) == 5
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# 审阅回归：PATCH tags 语义 / 向量-块一致性 / 过滤漏召回 / rebuild 语义与回滚
|
||||
# --------------------------------------------------------------------------- #
|
||||
@@ -563,9 +640,10 @@ def test_rebuild_failure_restores_old_index(vault, monkeypatch) -> None:
|
||||
assert repository.stats() == before # 旧索引已恢复，无半成品
|
||||
|
||||
|
||||
def test_first_rebuild_failure_removes_partial_database(vault, monkeypatch) -> None:
|
||||
def test_first_rebuild_failure_leaves_no_partial_index(vault, monkeypatch) -> None:
|
||||
"""首次启动没有旧库时，失败也不能留下已经写入的部分索引。"""
|
||||
from app.services import index_service
|
||||
from app import repository
|
||||
|
||||
_write_vault(
|
||||
vault,
|
||||
@@ -574,17 +652,17 @@ def test_first_rebuild_failure_removes_partial_database(vault, monkeypatch) -> N
|
||||
real_index = index_service.index_note
|
||||
calls = {"count": 0}
|
||||
|
||||
async def fail_on_second(parsed):
|
||||
async def fail_on_second(parsed, **kwargs):
|
||||
calls["count"] += 1
|
||||
if calls["count"] == 2:
|
||||
raise RuntimeError("injected first-rebuild failure")
|
||||
await real_index(parsed)
|
||||
await real_index(parsed, **kwargs)
|
||||
|
||||
monkeypatch.setattr(index_service, "index_note", fail_on_second)
|
||||
with pytest.raises(RuntimeError):
|
||||
asyncio.run(index_service.rebuild(IndexRebuildRequest(scope="all")))
|
||||
|
||||
assert not get_settings().db_path.exists()
|
||||
assert repository.stats() == {"notes": 0, "blocks": 0}
|
||||
|
||||
|
||||
def test_rebuild_preserves_task_note_links(vault) -> None:
|
||||
|
||||
@@ -66,6 +66,67 @@ async def seed():
|
||||
return apple, banana
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["cancel", "write"])
|
||||
def test_rebuild_failure_preserves_concurrent_configuration_and_all_indexes(runtime, monkeypatch, failure):
|
||||
from app.container import container
|
||||
from app.contracts import ModelRoutingConfig, ProviderConfig, ProviderType
|
||||
from app.services import task_service
|
||||
|
||||
async def scenario():
|
||||
apple, _ = await seed()
|
||||
task = task_service.create_task(title="before", note_id=apple.note_id)
|
||||
before = {table: [tuple(row) for row in rows(f"SELECT * FROM {table}")]
|
||||
for table in ("notes", "blocks", "blocks_fts", "vec_blocks", "index_meta", "routed_block_vectors")}
|
||||
container.model_routing.update(ModelRoutingConfig())
|
||||
entered, release = asyncio.Event(), asyncio.Event()
|
||||
original_embed = runtime.embed
|
||||
|
||||
async def pending_embed(texts):
|
||||
entered.set()
|
||||
await release.wait()
|
||||
return await original_embed(texts)
|
||||
|
||||
monkeypatch.setattr(runtime, "embed", pending_embed)
|
||||
original_index = index_service.index_note
|
||||
writes = 0
|
||||
|
||||
async def fail_write(parsed, **kwargs):
|
||||
nonlocal writes
|
||||
await original_index(parsed, **kwargs)
|
||||
writes += 1
|
||||
if writes == 2:
|
||||
raise RuntimeError("injected write failure")
|
||||
|
||||
if failure == "write":
|
||||
monkeypatch.setattr(index_service, "index_note", fail_write)
|
||||
rebuilding = asyncio.create_task(index_service.rebuild(IndexRebuildRequest()))
|
||||
await asyncio.wait_for(entered.wait(), timeout=5)
|
||||
saved = container.model_routing.update(container.model_routing.configuration())
|
||||
config = ProviderConfig(provider_id="concurrent", provider_type=ProviderType.openai_compatible,
|
||||
name="saved during rebuild", base_url="https://unused.invalid/v1")
|
||||
container.providers.register(config, container.provider_factory.build(config))
|
||||
task_service.update_task(task.task_id, {"title": "saved during rebuild"})
|
||||
# Preparation keeps the old searchable index intact while API I/O is pending.
|
||||
assert repository.stats()["notes"] == 2
|
||||
if failure == "cancel":
|
||||
rebuilding.cancel()
|
||||
expected = asyncio.CancelledError
|
||||
else:
|
||||
release.set()
|
||||
expected = RuntimeError
|
||||
with pytest.raises(expected):
|
||||
await rebuilding
|
||||
assert container.model_routing.configuration().version == saved.config.version
|
||||
assert rows("SELECT provider_id FROM provider_configs")[-1][0] == "concurrent"
|
||||
restored = task_service.get_task(task.task_id)
|
||||
assert restored.title == "saved during rebuild"
|
||||
assert restored.note_id == apple.note_id
|
||||
for table, values in before.items():
|
||||
assert [tuple(row) for row in rows(f"SELECT * FROM {table}")] == values
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def local_engine():
|
||||
return RetrievalEngine(HashEmbeddingProvider(), LexicalReranker(), SqliteVecStore())
|
||||
|
||||
|
||||
Reference in New Issue
Block a user