fix(provider): 修复索引事务回滚与工具名分片并同步主分支

This commit is contained in:
2026-09-04 07:15:52 +08:00
30 changed files with 2232 additions and 98 deletions
+8
View File
@@ -0,0 +1,8 @@
"""Benchmark 服务:RAG / Agent 数据集注册、指标计算与运行管理。
模块划分:
- metrics.py 纯函数指标(Hit@K / Recall@K / MRR / CitationHit / 分位数)
- datasets.py 受控目录的 Dataset 注册与校验
- rag.py RAG Benchmark Runner(调用 retrieval.engine.search
- service.py 运行注册表、配置快照与报告组装
"""
+198
View File
@@ -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},
)
+58
View File
@@ -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
+158
View File
@@ -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,
)
+348
View File
@@ -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)
+4
View File
@@ -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"))
),
)
+154
View File
@@ -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
+6 -9
View File
@@ -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
View File
@@ -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]:
+2
View File
@@ -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]]:
+67 -16
View File
@@ -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: # hybridRRF 融合
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) / spannorm >= 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,
+2
View File
@@ -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
+13 -3
View File
@@ -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()
+177
View File
@@ -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
+24 -27
View File
@@ -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)
+25 -9
View File
@@ -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