DeterministicEmbedder(可换 BGE-M3 的接口位)+ kb_chunk 双索引写入 + 硬过滤→向量→关键词→RRF 重排管线 + kb_gap 缺口埋点(T-P2-03/04/06,TD-08)。 实现: - embedding.py:EmbeddingProvider Protocol + DeterministicEmbedder (sha256 种子→归一化 1024 维假向量,同文本同向量/异文本近似正交) + build_embedding_provider 工厂(未知模型名抛 CONFIG_ERROR)。 - repo.py:upsert_chunks(整实体先删后插,逐行带 embedding_model/version, 可重跑判据 T-P2-03);hybrid_search(§8.3 CTE:硬过滤 published/ entity_types/occasion/avoid_body_shapes → <=> 余弦 topN → ts_rank topN → RRF);record_gap(咨询锁串行化同 query+filter_ctx 的 hit_count 递增)。 - rerank.py:RRF/权重/新鲜度纯数学(独立于 SQL 以便单测逐项断言)。 - service.py:RetrievalService.search 编排 embed→检索→(空结果)gap 埋点, SearchOutcome.degraded 标记降级。 测试:单测 +30(test_embedding 确定性/维度/归一化/工厂拒绝、test_rerank 公式逐项断言、test_retrieval_service 编排打桩);集成 +5 (test_hybrid_retrieval:禁忌实体同文最近向量仍零出现、draft 零出现、 关键词路径、换模型整组替换、kb_gap hit_count=2、partial 索引 EXPLAIN 验证)。 全绿:ruff / mypy(retrieval+shared+domain) / lint-imports / 268 单测 / 21 集成。 偏差说明: 1) 新增 rerank.py(任务文档只点名 embedding/repo/service,但重排数学必须 不连库可单测,独立模块是最低成本实现,仍属 packages/retrieval/**)。 2) upsert_chunks 的 upsert 粒度是实体(先删后插):kb_chunk 无业务自然键, 行级 ON CONFLICT 不可行;换模型名=新版本整组替换,不留旧版本残向量。 3) 禁忌过滤用 COALESCE(attrs->avoid_body_shapes, '[]')兜底:jsonb NULL 的 ? 运算结果是 NULL,直接 NOT 会把未声明禁忌的好实体一并排除。 4) 关键词路过滤用 tsv @@ tsquery 而非 ts_rank>0:ts_rank 对不匹配文本返回 ~1e-20 正 epsilon,阈值过滤会把无关 chunk 全放进关键词路。 5) RRF 分量按 1/(k+1) 归一到 (0,1],使 0.45/0.2/0.25/0.1 权重量纲可比 (原始倒数 0.45×1/61 与规则项 0.25×1 完全不可比)。
180 lines
5.8 KiB
Python
180 lines
5.8 KiB
Python
"""RetrievalService 编排逻辑:repo 层打桩,不连库(编排正确性与 DB 无关)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from packages.retrieval import repo
|
|
from packages.retrieval import service as service_mod
|
|
from packages.retrieval.rerank import ScoredChunk
|
|
from packages.retrieval.service import RetrievalService, UserContext
|
|
|
|
|
|
class FakeEmbedder:
|
|
"""满足 EmbeddingProvider Protocol 的最小桩:记录调用,返回固定向量."""
|
|
|
|
model_name = "fake"
|
|
version = "v1"
|
|
dimensions = 4
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[list[str]] = []
|
|
|
|
async def embed(self, texts: list[str]) -> list[list[float]]:
|
|
self.calls.append(list(texts))
|
|
return [[0.25, 0.25, 0.25, 0.25] for _ in texts]
|
|
|
|
|
|
class FakeSession:
|
|
def __init__(self) -> None:
|
|
self.commits = 0
|
|
self.rollbacks = 0
|
|
|
|
async def commit(self) -> None:
|
|
self.commits += 1
|
|
|
|
async def rollback(self) -> None:
|
|
self.rollbacks += 1
|
|
|
|
|
|
def _scored(chunk_id: int = 1) -> ScoredChunk:
|
|
return ScoredChunk(
|
|
chunk_id=chunk_id,
|
|
entity_id=100,
|
|
chunk_type="text",
|
|
text="chunk",
|
|
vec_score=1.0,
|
|
kw_score=0.0,
|
|
rule_score=0.0,
|
|
freshness_score=1.0,
|
|
final_score=0.55,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def stubbed_repo(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]:
|
|
"""把 repo 的两个 IO 函数换成记录桩,service 逻辑照常执行."""
|
|
|
|
calls: dict[str, Any] = {"hybrid_search": None, "gaps": [], "results": []}
|
|
|
|
async def fake_hybrid_search(
|
|
session: Any, query_vec: Any, query_text: str, **kwargs: Any
|
|
) -> list[ScoredChunk]:
|
|
calls["hybrid_search"] = {"query_vec": query_vec, "query_text": query_text, **kwargs}
|
|
return list(calls["results"])
|
|
|
|
async def fake_record_gap(session: Any, query_text: str, filter_ctx: Any) -> None:
|
|
calls["gaps"].append((query_text, filter_ctx))
|
|
|
|
monkeypatch.setattr(repo, "hybrid_search", fake_hybrid_search)
|
|
monkeypatch.setattr(repo, "record_gap", fake_record_gap)
|
|
return calls
|
|
|
|
|
|
async def test_search_with_hits_returns_results_without_gap(stubbed_repo: dict[str, Any]) -> None:
|
|
stubbed_repo["results"] = [_scored(1), _scored(2)]
|
|
embedder = FakeEmbedder()
|
|
session = FakeSession()
|
|
svc = RetrievalService(session, embedding=embedder)
|
|
|
|
outcome = await svc.search("面试穿什么", user_ctx=UserContext(body_shape="apple"))
|
|
|
|
assert outcome.degraded is False
|
|
assert [c.chunk_id for c in outcome.results] == [1, 2]
|
|
assert stubbed_repo["gaps"] == []
|
|
assert session.commits == 0 # 有结果的路径不写库,事务留给调用方
|
|
assert embedder.calls == [["面试穿什么"]] # embed 的输入就是 query 原文
|
|
|
|
|
|
async def test_search_forwards_user_context_to_hybrid(stubbed_repo: dict[str, Any]) -> None:
|
|
stubbed_repo["results"] = [_scored()]
|
|
svc = RetrievalService(FakeSession(), embedding=FakeEmbedder())
|
|
|
|
await svc.search(
|
|
"婚礼穿什么",
|
|
user_ctx=UserContext(
|
|
body_shape="pear",
|
|
occasion="wedding",
|
|
entity_types=("item", "style"),
|
|
rule_boost={42: 0.8},
|
|
),
|
|
)
|
|
|
|
seen = stubbed_repo["hybrid_search"]
|
|
assert seen is not None
|
|
assert seen["query_text"] == "婚礼穿什么"
|
|
assert seen["body_shape"] == "pear"
|
|
assert seen["occasion"] == "wedding"
|
|
assert seen["entity_types"] == ("item", "style")
|
|
assert seen["rule_boost"] == {42: 0.8}
|
|
|
|
|
|
async def test_empty_entity_types_is_forwarded_as_none(stubbed_repo: dict[str, Any]) -> None:
|
|
stubbed_repo["results"] = []
|
|
svc = RetrievalService(FakeSession(), embedding=FakeEmbedder())
|
|
|
|
await svc.search("随便看看", user_ctx=UserContext())
|
|
|
|
seen = stubbed_repo["hybrid_search"]
|
|
assert seen is not None
|
|
assert seen["entity_types"] is None # 空元组 = 不限类型,不是匹配空集合
|
|
|
|
|
|
async def test_no_results_records_gap_and_degrades(stubbed_repo: dict[str, Any]) -> None:
|
|
stubbed_repo["results"] = []
|
|
session = FakeSession()
|
|
svc = RetrievalService(session, embedding=FakeEmbedder())
|
|
|
|
outcome = await svc.search(
|
|
"不存在的知识点", user_ctx=UserContext(body_shape="apple", entity_types=("item",))
|
|
)
|
|
|
|
assert outcome.degraded is True
|
|
assert outcome.results == []
|
|
assert session.commits == 1 # gap 埋点立即提交:检索本身只读,别拖住调用方事务
|
|
assert stubbed_repo["gaps"] == [
|
|
(
|
|
"不存在的知识点",
|
|
{"body_shape": "apple", "occasion": None, "entity_types": ["item"]},
|
|
)
|
|
]
|
|
|
|
|
|
async def test_gap_failure_degrades_but_does_not_raise(
|
|
stubbed_repo: dict[str, Any], monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""埋点失败不该把"检索空结果"变成异常 —— 空结果本身是合法结局."""
|
|
|
|
async def broken_gap(session: Any, query_text: str, filter_ctx: Any) -> None:
|
|
raise RuntimeError("db down")
|
|
|
|
monkeypatch.setattr(repo, "record_gap", broken_gap)
|
|
stubbed_repo["results"] = []
|
|
session = FakeSession()
|
|
svc = RetrievalService(session, embedding=FakeEmbedder())
|
|
|
|
outcome = await svc.search("anything")
|
|
|
|
assert outcome.degraded is True
|
|
assert session.rollbacks == 1
|
|
|
|
|
|
async def test_default_context_when_user_ctx_omitted(stubbed_repo: dict[str, Any]) -> None:
|
|
stubbed_repo["results"] = []
|
|
svc = RetrievalService(FakeSession(), embedding=FakeEmbedder())
|
|
|
|
outcome = await svc.search("no ctx")
|
|
|
|
assert outcome.degraded is True
|
|
assert stubbed_repo["gaps"] == [
|
|
("no ctx", {"body_shape": None, "occasion": None, "entity_types": []})
|
|
]
|
|
|
|
|
|
def test_service_module_reexports_for_dependency_injection() -> None:
|
|
"""service 与 repo 共享同一模块对象:monkeypatch 打桩路径就是生产调用路径."""
|
|
|
|
assert service_mod.repo is repo
|