kplam
6ae91679e2
- .env.example 嵌入默认改为本地 BGE-M3（EMBEDDING_PROVIDER=local, dim=1024），NAS 全新部署即用 - 新增 app/services/reranker.py：基于 Ollama /api/rerank 的 cross-encoder 重排服务 - retriever 在 chunk 候选阶段接入语义精排，调用失败优雅降级为 RRF 顺序 - docker-compose 拉取 qwen3-reranker:0.6b，OLLAMA_MAX_LOADED_MODELS 提至 3 - .gitignore 排除 .workbuddy/ 与 .trae/（防止误提交项目记忆与 IDE spec） - 新增 reranker 单元与集成测试，全量测试 529 passed
414 lines
16 KiB
Python
414 lines
16 KiB
Python
"""分层检索引擎与检索 API 的单元测试（全部 mock，不联网、不起 Docker）"""
|
||
|
||
from typing import Any
|
||
|
||
import pytest
|
||
from fastapi.testclient import TestClient
|
||
from qdrant_client import models
|
||
|
||
from app.config import settings
|
||
from app.core.query_parser import ParsedQuery, RouteDecision
|
||
from app.core.retriever import Retriever
|
||
from app.core.sparse import SparseEncoder
|
||
from app.main import app
|
||
from app.models.search import SearchHit, SearchRequest, SearchResponse
|
||
from app.services.qdrant import COLLECTION_CHUNKS, COLLECTION_L1, COLLECTION_L2, COLLECTION_L3
|
||
|
||
|
||
def _point(
|
||
pid: str, doc_id: str = "", section_path: str | None = None, score: float = 1.0, **extra: Any
|
||
) -> models.ScoredPoint:
|
||
"""构造测试用 ScoredPoint，payload 模拟 chunk/节点结构"""
|
||
payload: dict[str, Any] = {"doc_id": doc_id, "text": f"text-{pid}", "title": f"title-{doc_id}", **extra}
|
||
if section_path is not None:
|
||
payload["section_path"] = section_path
|
||
return models.ScoredPoint(id=pid, version=0, score=score, payload=payload, vector=None)
|
||
|
||
|
||
class FakeQdrant:
|
||
"""按集合 + 调用顺序返回预置结果的假 QdrantService，记录每次调用的参数"""
|
||
|
||
def __init__(self, results: dict[str, list[list[models.ScoredPoint]]]) -> None:
|
||
self._results = results
|
||
self.calls: list[dict[str, Any]] = []
|
||
|
||
def _next(self, collection: str) -> list[models.ScoredPoint]:
|
||
queue = self._results.get(collection, [])
|
||
return queue.pop(0) if queue else []
|
||
|
||
async def search_dense(
|
||
self, collection: str, vector: list[float], limit: int, query_filter: models.Filter | None = None
|
||
) -> list[models.ScoredPoint]:
|
||
self.calls.append({"collection": collection, "method": "dense", "limit": limit, "filter": query_filter})
|
||
return self._next(collection)
|
||
|
||
async def search_hybrid(
|
||
self,
|
||
collection: str,
|
||
dense_vector: list[float],
|
||
sparse: tuple[list[int], list[float]],
|
||
limit: int,
|
||
query_filter: models.Filter | None = None,
|
||
) -> list[models.ScoredPoint]:
|
||
self.calls.append({"collection": collection, "method": "hybrid", "limit": limit, "filter": query_filter})
|
||
return self._next(collection)
|
||
|
||
|
||
class FakeParser:
|
||
"""返回固定路由决策的假 QueryParser"""
|
||
|
||
def __init__(self, route: RouteDecision) -> None:
|
||
self.route = route
|
||
|
||
async def parse_and_route(self, query: str) -> RouteDecision:
|
||
return self.route
|
||
|
||
|
||
class FakeEmbedding:
|
||
"""返回固定向量的假 EmbeddingService"""
|
||
|
||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||
return [[0.1, 0.2, 0.3] for _ in texts]
|
||
|
||
|
||
def _route(
|
||
fallback: bool = False, categories: tuple[str, ...] = ("技术文档",), rewrite: str = "rewrite query"
|
||
) -> RouteDecision:
|
||
return RouteDecision(
|
||
fallback=fallback,
|
||
filter_categories=None if fallback else list(categories),
|
||
reason="low_confidence" if fallback else "routed",
|
||
parsed=ParsedQuery(raw_query="q", rewrite=rewrite),
|
||
)
|
||
|
||
|
||
def _make_retriever(qdrant: FakeQdrant, route: RouteDecision, reranker: Any = None) -> Retriever:
|
||
return Retriever(
|
||
qdrant=qdrant, # type: ignore[arg-type]
|
||
query_parser=FakeParser(route), # type: ignore[arg-type]
|
||
embedding=FakeEmbedding(),
|
||
sparse_encoder=SparseEncoder(),
|
||
reranker=reranker,
|
||
)
|
||
|
||
|
||
class FakeReranker:
|
||
"""假 RerankerService：返回预设分数或抛出预设异常"""
|
||
|
||
def __init__(self, scores: list[float] | None = None, exc: Exception | None = None) -> None:
|
||
self.scores = scores if scores is not None else []
|
||
self.exc = exc
|
||
|
||
async def rerank(self, query: str, documents: list[str]) -> list[float]:
|
||
if self.exc is not None:
|
||
raise self.exc
|
||
return self.scores
|
||
|
||
|
||
def _calls(qdrant: FakeQdrant, collection: str) -> list[dict[str, Any]]:
|
||
return [c for c in qdrant.calls if c["collection"] == collection]
|
||
|
||
|
||
def _must_match_any(flt: models.Filter | None, key: str) -> list[str] | None:
|
||
"""提取 must 中指定 key 的 MatchAny 值，不存在返回 None"""
|
||
if flt is None:
|
||
return None
|
||
for cond in flt.must or []:
|
||
if cond.key == key:
|
||
return list(cond.match.any)
|
||
return None
|
||
|
||
|
||
def _filter_categories(flt: models.Filter | None) -> list[str] | None:
|
||
"""提取 min_should 中 category 条件的 MatchAny 值，不存在返回 None"""
|
||
if flt is None or flt.min_should is None:
|
||
return None
|
||
for cond in flt.min_should.conditions:
|
||
if cond.key == "category":
|
||
return list(cond.match.any)
|
||
return None
|
||
|
||
|
||
class TestRetriever:
|
||
"""分层检索流程：各层 filter 参数与回退路径"""
|
||
|
||
async def test_full_pipeline(self):
|
||
"""正常三级逐层：L1 候选 → L2 收窄 → L3 两路（含 2.5 级文档 b 路）→ chunk"""
|
||
qdrant = FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[_point("l1a", "d1"), _point("l1b", "d2")]],
|
||
COLLECTION_L2: [[_point("l2a", "d1", "章节A")]],
|
||
COLLECTION_L3: [
|
||
[_point("l3a", "d1", "章节A")], # a 路：L2 命中文档
|
||
[_point("l3b", "d2", "")], # b 路：2.5 级文档（无 L2 节点）
|
||
],
|
||
COLLECTION_CHUNKS: [
|
||
[
|
||
_point("c1", "d1", "章节A", doc_summary="总结1"),
|
||
_point("c2", "d2", "", doc_summary="总结2"),
|
||
]
|
||
],
|
||
}
|
||
)
|
||
retriever = _make_retriever(qdrant, _route())
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
# L1：categories 过滤、无 doc 限制，limit=l1_doc_top_n
|
||
l1_call = _calls(qdrant, COLLECTION_L1)[0]
|
||
assert l1_call["limit"] == settings.l1_doc_top_n
|
||
assert _filter_categories(l1_call["filter"]) == ["技术文档"]
|
||
assert _must_match_any(l1_call["filter"], "doc_id") is None
|
||
|
||
# L2：doc_ids 为 L1 候选（保序），limit=l2_section_top_n*候选数
|
||
l2_call = _calls(qdrant, COLLECTION_L2)[0]
|
||
assert l2_call["method"] == "dense"
|
||
assert l2_call["limit"] == settings.l2_section_top_n * 2
|
||
assert _must_match_any(l2_call["filter"], "doc_id") == ["d1", "d2"]
|
||
assert _filter_categories(l2_call["filter"]) == ["技术文档"]
|
||
|
||
# L3 两路：a 路 d1 + section「章节A」；b 路 d2（2.5 级文档）仅 doc 过滤
|
||
l3_calls = _calls(qdrant, COLLECTION_L3)
|
||
assert len(l3_calls) == 2
|
||
assert all(c["limit"] == settings.l3_top_n for c in l3_calls)
|
||
assert _must_match_any(l3_calls[0]["filter"], "doc_id") == ["d1"]
|
||
assert _must_match_any(l3_calls[0]["filter"], "section_path") == ["章节A"]
|
||
assert _must_match_any(l3_calls[1]["filter"], "doc_id") == ["d2"]
|
||
assert _must_match_any(l3_calls[1]["filter"], "section_path") is None
|
||
|
||
# chunk：doc_ids 为 L3 命中文档，section_paths 仅非空值
|
||
chunk_call = _calls(qdrant, COLLECTION_CHUNKS)[0]
|
||
assert chunk_call["limit"] == settings.retrieval_top_k
|
||
assert sorted(_must_match_any(chunk_call["filter"], "doc_id") or []) == ["d1", "d2"]
|
||
assert _must_match_any(chunk_call["filter"], "section_path") == ["章节A"]
|
||
|
||
# 响应组装
|
||
assert resp.query == "测试查询"
|
||
assert resp.fallback is False
|
||
assert resp.routed_categories == ["技术文档"]
|
||
assert [h.doc_id for h in resp.hits] == ["d1", "d2"]
|
||
assert resp.hits[0].text == "text-c1"
|
||
assert resp.hits[0].title == "title-d1"
|
||
assert resp.hits[0].section_path == "章节A"
|
||
assert resp.hits[0].doc_summary == "总结1"
|
||
assert resp.hits[0].score > 0
|
||
|
||
async def test_l2_empty_all_docs_go_l3_b_path(self):
|
||
"""L2 整体无命中：全部候选文档回退为 L3 单路 doc 级查询"""
|
||
qdrant = FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[_point("l1a", "d1"), _point("l1b", "d2")]],
|
||
COLLECTION_L2: [[]],
|
||
COLLECTION_L3: [[_point("l3a", "d1", "章节A")]],
|
||
COLLECTION_CHUNKS: [[_point("c1", "d1", "章节A")]],
|
||
}
|
||
)
|
||
retriever = _make_retriever(qdrant, _route())
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
l3_calls = _calls(qdrant, COLLECTION_L3)
|
||
assert len(l3_calls) == 1
|
||
assert _must_match_any(l3_calls[0]["filter"], "doc_id") == ["d1", "d2"]
|
||
assert _must_match_any(l3_calls[0]["filter"], "section_path") is None
|
||
assert [h.doc_id for h in resp.hits] == ["d1"]
|
||
|
||
async def test_l3_empty_fallback_doc_level_chunks(self):
|
||
"""L3 两路均无命中：chunk 层回退为 L1 候选文档级检索（无 section 过滤）"""
|
||
qdrant = FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[_point("l1a", "d1"), _point("l1b", "d2")]],
|
||
COLLECTION_L2: [[_point("l2a", "d1", "章节A")]],
|
||
COLLECTION_L3: [[], []],
|
||
COLLECTION_CHUNKS: [[_point("c1", "d2", "")]],
|
||
}
|
||
)
|
||
retriever = _make_retriever(qdrant, _route())
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
chunk_call = _calls(qdrant, COLLECTION_CHUNKS)[0]
|
||
assert _must_match_any(chunk_call["filter"], "doc_id") == ["d1", "d2"]
|
||
assert _must_match_any(chunk_call["filter"], "section_path") is None
|
||
assert [h.doc_id for h in resp.hits] == ["d2"]
|
||
|
||
async def test_l1_empty_global_chunk_fallback(self):
|
||
"""L1 无候选文档：直接全库 chunk 兜底（无 filter），fallback=True"""
|
||
qdrant = FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[]],
|
||
COLLECTION_CHUNKS: [[_point("c1", "d9", "章节X", doc_summary="总结9")]],
|
||
}
|
||
)
|
||
retriever = _make_retriever(qdrant, _route())
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
# 不再触发 L2/L3 查询
|
||
assert _calls(qdrant, COLLECTION_L2) == []
|
||
assert _calls(qdrant, COLLECTION_L3) == []
|
||
# 全库 chunk 检索：无 filter
|
||
chunk_calls = _calls(qdrant, COLLECTION_CHUNKS)
|
||
assert len(chunk_calls) == 1
|
||
assert chunk_calls[0]["filter"] is None
|
||
assert chunk_calls[0]["limit"] == settings.retrieval_top_k
|
||
assert resp.fallback is True
|
||
assert [h.doc_id for h in resp.hits] == ["d9"]
|
||
|
||
async def test_route_fallback_no_category_filter(self):
|
||
"""路由兜底时 categories=None，各层均不做类目过滤"""
|
||
qdrant = FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[_point("l1a", "d1")]],
|
||
COLLECTION_L2: [[]],
|
||
COLLECTION_L3: [[_point("l3a", "d1", "")]],
|
||
COLLECTION_CHUNKS: [[_point("c1", "d1", "")]],
|
||
}
|
||
)
|
||
retriever = _make_retriever(qdrant, _route(fallback=True))
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
# L1：categories 与 doc_ids 均空 → filter 为 None
|
||
assert _calls(qdrant, COLLECTION_L1)[0]["filter"] is None
|
||
# L2/L3/chunk：仅有 doc 级 must 条件，无类目 min_should
|
||
for collection in (COLLECTION_L2, COLLECTION_L3, COLLECTION_CHUNKS):
|
||
for call in _calls(qdrant, collection):
|
||
assert _filter_categories(call["filter"]) is None
|
||
assert resp.fallback is True
|
||
assert resp.routed_categories == []
|
||
|
||
async def test_top_k_override(self):
|
||
"""request.top_k 优先于 settings.retrieval_final_k"""
|
||
qdrant = FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[_point("l1a", "d1")]],
|
||
COLLECTION_L2: [[]],
|
||
COLLECTION_L3: [[_point("l3a", "d1", "")]],
|
||
COLLECTION_CHUNKS: [[_point(f"c{i}", "d1", "") for i in range(5)]],
|
||
}
|
||
)
|
||
retriever = _make_retriever(qdrant, _route())
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询", top_k=2))
|
||
assert len(resp.hits) == 2
|
||
|
||
|
||
class _FakeRetriever:
|
||
"""API 测试用假 Retriever：返回固定响应或抛异常"""
|
||
|
||
def __init__(self, response: SearchResponse | None = None, exc: Exception | None = None) -> None:
|
||
self.response = response
|
||
self.exc = exc
|
||
|
||
async def search(self, request: SearchRequest) -> SearchResponse:
|
||
if self.exc is not None:
|
||
raise self.exc
|
||
assert self.response is not None
|
||
return self.response
|
||
|
||
|
||
class TestSearchApi:
|
||
"""检索 API 统一响应包装"""
|
||
|
||
def test_search_ok(self, monkeypatch: pytest.MonkeyPatch):
|
||
"""正常检索 → {"code": 0, "data": ..., "message": "ok"}"""
|
||
response = SearchResponse(
|
||
query="q",
|
||
hits=[SearchHit(text="t", doc_id="d1", title="标题", section_path="", score=0.5, doc_summary="s")],
|
||
routed_categories=["技术文档"],
|
||
fallback=False,
|
||
)
|
||
monkeypatch.setattr("app.api.v1.search._retriever", _FakeRetriever(response=response))
|
||
|
||
client = TestClient(app)
|
||
resp = client.post("/api/v1/search", json={"query": "q"})
|
||
|
||
assert resp.status_code == 200
|
||
body = resp.json()
|
||
assert body["code"] == 0
|
||
assert body["message"] == "ok"
|
||
assert body["data"]["query"] == "q"
|
||
assert body["data"]["hits"][0]["doc_id"] == "d1"
|
||
assert body["data"]["routed_categories"] == ["技术文档"]
|
||
assert body["data"]["fallback"] is False
|
||
|
||
def test_search_error(self, monkeypatch: pytest.MonkeyPatch):
|
||
"""检索内部异常 → code 2000"""
|
||
monkeypatch.setattr("app.api.v1.search._retriever", _FakeRetriever(exc=RuntimeError("boom")))
|
||
|
||
client = TestClient(app)
|
||
resp = client.post("/api/v1/search", json={"query": "q"})
|
||
|
||
assert resp.status_code == 200
|
||
body = resp.json()
|
||
assert body["code"] == 2000
|
||
assert body["data"] is None
|
||
assert "boom" in body["message"]
|
||
|
||
def test_validation_error(self):
|
||
"""缺少必填 query 字段 → code 1001"""
|
||
client = TestClient(app)
|
||
resp = client.post("/api/v1/search", json={})
|
||
|
||
assert resp.status_code == 200
|
||
body = resp.json()
|
||
assert body["code"] == 1001
|
||
assert body["data"] is None
|
||
|
||
def test_health_kept(self):
|
||
"""现有健康检查接口不受影响"""
|
||
client = TestClient(app)
|
||
resp = client.get("/api/v1/health")
|
||
assert resp.status_code == 200
|
||
assert resp.json() == {"status": "ok"}
|
||
|
||
|
||
class TestReranker:
|
||
"""重排启用时对 chunk 候选精排；关闭或失败则退化为 RRF 顺序"""
|
||
|
||
@staticmethod
|
||
def _full_qdrant() -> FakeQdrant:
|
||
return FakeQdrant(
|
||
{
|
||
COLLECTION_L1: [[_point("l1a", "d1")]],
|
||
COLLECTION_L2: [[_point("l2a", "d1", "章节A")]],
|
||
COLLECTION_L3: [[_point("l3a", "d1", "章节A")]],
|
||
COLLECTION_CHUNKS: [
|
||
[
|
||
_point("c1", "d1", "章节A"),
|
||
_point("c2", "d1", "章节A"),
|
||
_point("c3", "d1", "章节A"),
|
||
]
|
||
],
|
||
}
|
||
)
|
||
|
||
async def test_enabled_reorders_by_relevance(self):
|
||
"""重排按相关性分数降序重排，并将 score 写回相关性分数"""
|
||
reranker = FakeReranker(scores=[0.1, 0.9, 0.5]) # c1<c3<c2
|
||
retriever = _make_retriever(self._full_qdrant(), _route(), reranker=reranker)
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
# 重排后顺序为 c2, c3, c1
|
||
assert [h.text for h in resp.hits] == ["text-c2", "text-c3", "text-c1"]
|
||
assert resp.hits[0].score == 0.9
|
||
|
||
async def test_disabled_keeps_rrf_order(self):
|
||
"""重排未注入（关闭）时保持 RRF 融合原始顺序"""
|
||
retriever = _make_retriever(self._full_qdrant(), _route(), reranker=None)
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
assert [h.text for h in resp.hits] == ["text-c1", "text-c2", "text-c3"]
|
||
|
||
async def test_failure_falls_back_to_rrf(self):
|
||
"""重排调用异常时退化为 RRF 顺序，检索不中断"""
|
||
reranker = FakeReranker(exc=RuntimeError("ollama 不可用"))
|
||
retriever = _make_retriever(self._full_qdrant(), _route(), reranker=reranker)
|
||
|
||
resp = await retriever.search(SearchRequest(query="测试查询"))
|
||
|
||
assert [h.text for h in resp.hits] == ["text-c1", "text-c2", "text-c3"]
|