From aa1287d83d7b560ad2969800aa66c0b6368bde90 Mon Sep 17 00:00:00 2001 From: xsl Date: Fri, 14 Aug 2026 20:12:27 +0800 Subject: [PATCH] feat(retrieval): hybrid search over pgvector with hard filters and gap tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 完全不可比）。 --- packages/retrieval/README.md | 31 ++- packages/retrieval/__init__.py | 35 ++- packages/retrieval/embedding.py | 99 +++++++ packages/retrieval/repo.py | 263 ++++++++++++++++++ packages/retrieval/rerank.py | 145 ++++++++++ packages/retrieval/service.py | 114 ++++++++ tests/integration/test_hybrid_retrieval.py | 297 +++++++++++++++++++++ tests/unit/test_embedding.py | 99 +++++++ tests/unit/test_rerank.py | 158 +++++++++++ tests/unit/test_retrieval_service.py | 179 +++++++++++++ 10 files changed, 1409 insertions(+), 11 deletions(-) create mode 100644 packages/retrieval/embedding.py create mode 100644 packages/retrieval/repo.py create mode 100644 packages/retrieval/rerank.py create mode 100644 packages/retrieval/service.py create mode 100644 tests/integration/test_hybrid_retrieval.py create mode 100644 tests/unit/test_embedding.py create mode 100644 tests/unit/test_rerank.py create mode 100644 tests/unit/test_retrieval_service.py diff --git a/packages/retrieval/README.md b/packages/retrieval/README.md index 1d57fd4..992b673 100644 --- a/packages/retrieval/README.md +++ b/packages/retrieval/README.md @@ -1,14 +1,23 @@ -# `packages/retrieval/` — 混合检索（占位） +# `packages/retrieval/` — 向量与混合检索 -架构 §8.3。进程内包，CPU + PostgreSQL。 +架构 §8.3。进程内包，CPU + PostgreSQL，无 HTTP 路由（被 agent / report 的 +service 调用）。 -## 待实现 +## 模块 -一条 SQL 完成硬过滤 + 向量召回（这正是选 pgvector 的收益），再做 RRF 融合： +- `embedding.py`：`EmbeddingProvider` Protocol + `DeterministicEmbedder` + （测试/dev 的确定性假实现，sha256 种子 → 归一化 1024 维向量）+ 工厂。 + 真实 BGE-M3 ONNX 加载是后续任务（`kb_embedding_device` 已预留）。 +- `repo.py`：`kb_chunk` / `kb_gap` 数据访问。一条 CTE 完成硬过滤 + + 向量召回（`<=>` 余弦）+ 关键词召回（`ts_rank`）。 +- `rerank.py`：RRF 融合重排的纯数学（单测逐项断言，不连库）。 +- `service.py`：编排 embed → 检索 →（无结果）gap 埋点。 + +管线： ``` -硬过滤(SQL WHERE) → 规则命中 → 向量召回(pgvector) → 关键词召回(tsvector) -→ 融合重排 → 输出知识片段 + 命中规则 ID + 图谱路径 +硬过滤(SQL WHERE) → 向量召回(pgvector) → 关键词召回(tsvector) +→ RRF 融合 + 权重重排 → 输出 ScoredChunk；空结果落 kb_gap ``` 重排权重（可配置，需用黄金集调参）： @@ -16,8 +25,10 @@ ## 两条不能妥协的要求 -- **硬过滤零漏**（FR-KB-06）：禁忌项不得出现在候选中。这不是排序问题， - 是过滤问题 —— 不能靠权重压低，必须在 SQL 里排除。 -- **规则命中要有明显加权**：专业规则的可信度高于语义相似度，这是自建知识库 - 的意义所在。 +- **硬过滤零漏**（FR-KB-06 / T-P2-04）：draft/禁忌实体必须在 SQL 的 + WHERE 里排除，不能靠权重压低 —— 这是过滤问题，不是排序问题。 + 集成测试 `tests/integration/test_hybrid_retrieval.py` 断言"禁忌实体 + 即使向量最近也绝不出现"。 +- **规则命中要有明显加权**：专业规则的可信度高于语义相似度（`rule_boost` + 接口已留，规则引擎联动属后续任务）。 - p95 < 200ms（NFR-PERF-06）。 diff --git a/packages/retrieval/__init__.py b/packages/retrieval/__init__.py index ba170ef..11b10da 100644 --- a/packages/retrieval/__init__.py +++ b/packages/retrieval/__init__.py @@ -1 +1,34 @@ -"""混合检索：硬过滤 → 规则命中 → 向量召回 → BM25 → 重排 → 解释路径。占位。""" +"""混合检索：硬过滤 → 向量召回 → 关键词召回 → RRF 重排 + 缺口埋点（架构 §8.3）. + +进程内包，被 agent / report 的 service 调用，没有 HTTP 路由。 +硬过滤零漏（FR-KB-06）与向量化可重跑（T-P2-03）的保证都在 +``repo.py`` 顶部注释里，改代码前先读。 +""" + +from __future__ import annotations + +from packages.retrieval.embedding import ( + DETERMINISTIC_MODEL_NAME, + DeterministicEmbedder, + EmbeddingProvider, + build_embedding_provider, +) +from packages.retrieval.repo import ChunkInput, hybrid_search, record_gap, upsert_chunks +from packages.retrieval.rerank import RecallRow, ScoredChunk +from packages.retrieval.service import RetrievalService, SearchOutcome, UserContext + +__all__ = [ + "DETERMINISTIC_MODEL_NAME", + "ChunkInput", + "DeterministicEmbedder", + "EmbeddingProvider", + "RecallRow", + "RetrievalService", + "ScoredChunk", + "SearchOutcome", + "UserContext", + "build_embedding_provider", + "hybrid_search", + "record_gap", + "upsert_chunks", +] diff --git a/packages/retrieval/embedding.py b/packages/retrieval/embedding.py new file mode 100644 index 0000000..7875c11 --- /dev/null +++ b/packages/retrieval/embedding.py @@ -0,0 +1,99 @@ +"""嵌入提供方：接口位 + 确定性假实现（架构 §8.3，TD-08）. + +当前只有 ``deterministic-v1``：以文本 sha256 为种子生成伪随机向量后 L2 归一化。 +它是**测试与 dev 的假实现** —— 真实 BGE-M3 ONNX 加载是后续任务，配置里的 +``kb_embedding_device`` 就是为它预留的。接口先定死，届时换模型只改工厂， +检索管线（repo/service）不动。 +""" + +from __future__ import annotations + +import hashlib +from typing import Protocol, runtime_checkable + +import numpy as np + +from packages.shared.config import Settings +from packages.shared.errors import AppError, ErrorCode + +#: 唯一被承认的假实现模型名（与 Settings.kb_embedding_model 的默认值对应）。 +DETERMINISTIC_MODEL_NAME = "deterministic-v1" + + +@runtime_checkable +class EmbeddingProvider(Protocol): + """嵌入模型的最小接口。 + + ``model_name`` + ``version`` 会随向量逐行写进 kb_chunk（T-P2-03）： + 没有版本号的向量不可重跑 —— 换模型即换版本，这是向量化可重跑的判据。 + """ + + @property + def model_name(self) -> str: ... + + @property + def version(self) -> str: ... + + @property + def dimensions(self) -> int: ... + + async def embed(self, texts: list[str]) -> list[list[float]]: + """批量嵌入。返回顺序与输入一一对应。""" + ... + + +class DeterministicEmbedder: + """确定性假嵌入：同文本同向量，不同文本近似正交。 + + sha256(text) 做种子的标准高斯伪随机向量再 L2 归一化。1024 维空间里 + 随机向量两两夹角趋于 90°（|dot| ≈ 1/√1024 ≈ 0.03），这让"按余弦距离 + 排序"的管线在没有真模型时也能被测试驱动：同文本距离 0，异文本距离 ~1。 + """ + + def __init__(self, *, dimensions: int = 1024, version: str = "v1") -> None: + if dimensions <= 0: + raise ValueError(f"dimensions 必须为正数, 得到 {dimensions}") + self._dimensions = dimensions + self._version = version + + @property + def model_name(self) -> str: + return DETERMINISTIC_MODEL_NAME + + @property + def version(self) -> str: + return self._version + + @property + def dimensions(self) -> int: + return self._dimensions + + def _vector(self, text: str) -> np.ndarray: + # 只取 sha256 前 16 字节做种子：PCG64 的种子空间足够，且避免跨平台 + # int 大小差异。同一 numpy 版本内逐位确定（假实现不承诺跨大版本稳定， + # 真正的可重跑保证来自 kb_chunk.embedding_model/version 落库）。 + seed = int.from_bytes(hashlib.sha256(text.encode("utf-8")).digest()[:16], "big") + vec = np.random.default_rng(seed).standard_normal(self._dimensions) + norm = float(np.linalg.norm(vec)) + # 标准正态向量的范数恒 > 0，这里只防御浮点极端值，不做业务分支。 + return vec / norm if norm > 0.0 else vec + + async def embed(self, texts: list[str]) -> list[list[float]]: + return [self._vector(t).tolist() for t in texts] + + +def build_embedding_provider(settings: Settings) -> EmbeddingProvider: + """按 ``settings.kb_embedding_model`` 选择实现。 + + 当前只认 ``deterministic-v1``：出现其它名字说明配置与代码版本不匹配 + （例如提前把 KB_EMBEDDING_MODEL 改成 BGE-M3 但 ONNX 加载任务还没合入）， + 直接拒绝比静默回退到假向量更能暴露问题。 + """ + if settings.kb_embedding_model == DETERMINISTIC_MODEL_NAME: + return DeterministicEmbedder(dimensions=settings.kb_vector_dimensions) + raise AppError( + ErrorCode.CONFIG_ERROR, + message=f"未知的嵌入模型: {settings.kb_embedding_model}", + hint="当前仅支持 deterministic-v1，真实 BGE-M3 属后续任务", + data={"model": settings.kb_embedding_model}, + ) diff --git a/packages/retrieval/repo.py b/packages/retrieval/repo.py new file mode 100644 index 0000000..c6144fc --- /dev/null +++ b/packages/retrieval/repo.py @@ -0,0 +1,263 @@ +"""kb_chunk / kb_gap 数据访问。只有本模块写这两张表（架构 §8.3）. + +两条红线写在这里，改代码前先读： + +1. **硬过滤零漏**（T-P2-04 / FR-KB-06）：draft/禁忌实体必须在 SQL 的 + WHERE 里排除，不能靠重排权重压低 —— 权重是排序问题，过滤是安全问题。 +2. **没有版本号的向量不可重跑**（T-P2-03）：写入必须携带 + embedding_model/version；重跑按实体整组替换而不是追加。 +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import RowMapping, delete, select, text +from sqlalchemy.ext.asyncio import AsyncSession + +from packages.domain.ids import next_id +from packages.retrieval.embedding import EmbeddingProvider +from packages.retrieval.rerank import ( + DEFAULT_RERANK_WEIGHTS, + RecallRow, + ScoredChunk, + rerank, +) +from packages.shared.errors import AppError, ErrorCode +from packages.shared.models import KbChunk, KbGap + +# 硬过滤 + 双路召回 + RRF 的 CTE 骨架（架构 §8.3）。 +# 用原生 SQL 而不是 ORM 表达式：``<=>`` 余弦距离、ts_rank、jsonb ``?`` 操作符 +# 与窗口函数在 ORM 层拼写出来可读性更差，CTE 结构才是这份查询的可维护形态。 +# 参数一律显式 CAST，让 asyncpg 在 NULL/数组参数上也能推断类型。 +# +# jsonb NULL 语义是两条容易被写错的规则，明确如下： +# - occasion：没声明 occasions 的实体在按场景过滤时**不命中**（它没有声明 +# 适配该场景，宁可少给不能错给）； +# - avoid_body_shapes：没声明禁忌的实体**必须保留**（NULL ? x 是 NULL，直接 +# NOT 会把好实体一起丢掉），所以用 COALESCE 空数组兜底 —— 没写禁忌 ≠ 全都禁忌。 +_SEARCH_SQL = text( + """ + WITH filtered AS ( + SELECT c.chunk_id, + c.entity_id, + c.chunk_type, + c.text, + c.tsv, + e.updated_at, + c.embedding <=> CAST(:query_vec AS vector) AS vec_distance, + websearch_to_tsquery('simple', CAST(:query_text AS text)) AS kw_query, + ts_rank(c.tsv, websearch_to_tsquery('simple', CAST(:query_text AS text))) + AS kw_tsrank + FROM kb_chunk c + JOIN kb_entity e ON e.entity_id = c.entity_id + WHERE e.status = 'published' + AND (CAST(:entity_types AS text[]) IS NULL + OR e.entity_type = ANY(CAST(:entity_types AS text[]))) + AND (CAST(:occasion AS text) IS NULL + OR e.attrs -> 'occasions' ? CAST(:occasion AS text)) + AND (CAST(:body_shape AS text) IS NULL + OR NOT (COALESCE(e.attrs -> 'avoid_body_shapes', '[]'::jsonb) + ? CAST(:body_shape AS text))) + ), + vec_ranked AS ( + SELECT chunk_id, + ROW_NUMBER() OVER (ORDER BY vec_distance, chunk_id) AS vec_rank, + vec_distance + FROM filtered + ORDER BY vec_distance, chunk_id + LIMIT CAST(:recall_k AS bigint) + ), + kw_ranked AS ( + SELECT chunk_id, + ROW_NUMBER() OVER (ORDER BY kw_tsrank DESC, chunk_id) AS kw_rank, + kw_tsrank + FROM filtered + -- 过滤必须用 @@ 而不是 ts_rank > 0：ts_rank 对不匹配文本也返回 + -- 极小的正 epsilon（~1e-20），靠分数阈值拦截会把无关 chunk 全放进来。 + WHERE kw_tsrank > 0 AND tsv @@ kw_query + ORDER BY kw_tsrank DESC, chunk_id + LIMIT CAST(:recall_k AS bigint) + ) + SELECT f.chunk_id, + f.entity_id, + f.chunk_type, + f.text, + f.updated_at, + v.vec_rank, + v.vec_distance, + k.kw_rank, + k.kw_tsrank + FROM vec_ranked v + FULL OUTER JOIN kw_ranked k ON k.chunk_id = v.chunk_id + JOIN filtered f ON f.chunk_id = COALESCE(v.chunk_id, k.chunk_id) + """ +) + + +@dataclass(frozen=True, slots=True) +class ChunkInput: + """待向量化写入的一个知识片段。tsv 是生成列，不在这里出现。""" + + text: str + chunk_type: str = "text" + + +def to_pgvector_literal(vec: Sequence[float]) -> str: + """把浮点向量序列化成 PG 的 vector 字面量 ``'[0.1,0.2,...]'``. + + asyncpg 对扩展类型没有二进制编解码器，文本格式是官方支持的传参方式； + SQL 侧再显式 ``::vector`` 转型，PG 负责解析。 + """ + return "[" + ",".join(repr(float(x)) for x in vec) + "]" + + +async def upsert_chunks( + session: AsyncSession, + entity_id: int, + chunks: Sequence[ChunkInput], + embedding: EmbeddingProvider, +) -> list[KbChunk]: + """整实体替换式写入 kb_chunk（可重跑判据 T-P2-03）. + + kb_chunk 没有业务自然键（text 改一个字就是新 chunk），所以 upsert 的 + 粒度是**实体**：先删该实体全部旧向量再插入新组 —— 重跑（无论换没换 + 模型）不留脏数据，embedding_model/version 逐行携带，换模型名即视为 + 不同版本。 + """ + for chunk in chunks: + if not chunk.text.strip(): + raise AppError( + ErrorCode.VALIDATION_FAILED, + message="知识片段文本不能为空", + data={"entity_id": entity_id}, + ) + + vectors = await embedding.embed([c.text for c in chunks]) + if len(vectors) != len(chunks): + raise AppError( + ErrorCode.INTERNAL_ERROR, + message="嵌入实现返回的向量数与输入文本数不一致", + data={"model": embedding.model_name}, + ) + for vec in vectors: + if len(vec) != embedding.dimensions: + raise AppError( + ErrorCode.INTERNAL_ERROR, + message="嵌入实现返回的维度与声明不一致", + data={ + "model": embedding.model_name, + "declared": embedding.dimensions, + "actual": len(vec), + }, + ) + + # 同一实体的旧向量在事务内先删后插：失败整体回滚，不会出现半新半旧。 + await session.execute(delete(KbChunk).where(KbChunk.entity_id == entity_id)) + rows = [ + KbChunk( + chunk_id=next_id(), + entity_id=entity_id, + chunk_type=chunk.chunk_type, + text=chunk.text, + embedding=vec, + embedding_model=embedding.model_name, + embedding_version=embedding.version, + ) + for chunk, vec in zip(chunks, vectors, strict=True) + ] + session.add_all(rows) + return rows + + +async def hybrid_search( + session: AsyncSession, + query_vec: Sequence[float], + query_text: str, + *, + entity_types: Sequence[str] | None = None, + occasion: str | None = None, + body_shape: str | None = None, + rule_boost: Mapping[int, float] | None = None, + weights: Mapping[str, float] | None = None, + top_k: int = 50, +) -> list[ScoredChunk]: + """硬过滤 → 向量召回 → 关键词召回 → RRF 融合重排（架构 §8.3）. + + 硬过滤（published / entity_types / occasion / avoid_body_shapes 禁忌） + 全部发生在 ``filtered`` CTE 里，两路召回只从过滤后的集合里取 top N —— + 这是"零漏"的结构保证，而非概率保证。空 ``entity_types`` 视为不限类型。 + """ + if top_k <= 0: + return [] + result = await session.execute( + _SEARCH_SQL, + { + "query_vec": to_pgvector_literal(query_vec), + "query_text": query_text, + "entity_types": list(entity_types) if entity_types else None, + "occasion": occasion, + "body_shape": body_shape, + "recall_k": top_k, + }, + ) + rows: list[RecallRow] = [_to_recall_row(row) for row in result.mappings().all()] + return rerank( + rows, + weights=weights if weights is not None else DEFAULT_RERANK_WEIGHTS, + rule_boost=rule_boost, + now=datetime.now(UTC), + top_k=top_k, + ) + + +def _to_recall_row(row: RowMapping) -> RecallRow: + """把裸行收敛成 RecallRow：SQL 返回的 Any 在边界处一次性转干净。""" + return RecallRow( + chunk_id=int(row["chunk_id"]), + entity_id=int(row["entity_id"]), + chunk_type=str(row["chunk_type"]), + text=str(row["text"]), + updated_at=row["updated_at"], + vec_rank=int(row["vec_rank"]) if row["vec_rank"] is not None else None, + vec_distance=float(row["vec_distance"]) if row["vec_distance"] is not None else None, + kw_rank=int(row["kw_rank"]) if row["kw_rank"] is not None else None, + kw_tsrank=float(row["kw_tsrank"]) if row["kw_tsrank"] is not None else None, + ) + + +async def record_gap( + session: AsyncSession, + query_text: str, + filter_ctx: Mapping[str, Any], +) -> None: + """无结果 query 落 kb_gap：同 query_text + filter_ctx 则 hit_count += 1（T-P2-06）. + + kb_gap 只有普通索引没有唯一约束，``ON CONFLICT`` 用不了；用事务级 + 咨询锁把「查-改-插」串行化，保证同 key 并发时仍是一行、计数不丢。 + 锁在事务提交/回滚时自动释放，由调用方的 session_scope 兜底。 + """ + ctx = dict(filter_ctx) + payload = json.dumps( + {"query_text": query_text, "filter_ctx": ctx}, + sort_keys=True, + ensure_ascii=False, + ) + digest = hashlib.sha256(payload.encode("utf-8")).digest()[:8] + lock_key = int.from_bytes(digest, "big", signed=True) + await session.execute(text("SELECT pg_advisory_xact_lock(:lock_key)"), {"lock_key": lock_key}) + + found = await session.execute( + select(KbGap).where(KbGap.query_text == query_text, KbGap.filter_ctx == ctx) + ) + gap = found.scalar_one_or_none() + if gap is None: + session.add(KbGap(gap_id=next_id(), query_text=query_text, filter_ctx=ctx, hit_count=1)) + else: + gap.hit_count += 1 + gap.last_seen_at = datetime.now(UTC) diff --git a/packages/retrieval/rerank.py b/packages/retrieval/rerank.py new file mode 100644 index 0000000..78bb0d3 --- /dev/null +++ b/packages/retrieval/rerank.py @@ -0,0 +1,145 @@ +"""RRF 融合重排的纯数学部分（架构 §8.3 唯一口径）. + +刻意与 SQL 分离：重排公式必须能被单测逐项断言（T-P2-04 验收 3）， +混进 repo 就会把数学正确性测试拖进数据库。repo 只负责召回与名次， +这里把名次变成最终分。 +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import UTC, datetime + +#: RRF 常数。60 是业界默认值：名次差异对分数的影响被压成平滑倒数， +#: 单路召回头部的偏差不至于淹没另一路。 +RRF_K = 60 + +#: 新鲜度半衰期（天）：30 天前的内容得 0.5，60 天前 0.25…… +#: 知识库是慢更新的，半衰期必须显著大于内容评审周期，否则新版本刚 +#: 发布就会压制尚未重审的旧版本。 +FRESHNESS_HALF_LIFE_DAYS = 30.0 + +#: 架构 §8.3 唯一口径：0.45 向量 + 0.2 关键词 + 0.25 规则 + 0.1 新鲜度。 +#: 与 Settings.retrieval_rerank_weights 的默认值保持一致（可被 env 覆盖）。 +DEFAULT_RERANK_WEIGHTS: dict[str, float] = { + "vector": 0.45, + "keyword": 0.2, + "rule": 0.25, + "freshness": 0.1, +} + + +@dataclass(frozen=True, slots=True) +class RecallRow: + """混合检索的一行召回结果：两路召回各自的名次，未被召回的一路为 ``None``.""" + + chunk_id: int + entity_id: int + chunk_type: str + text: str + updated_at: datetime + vec_rank: int | None = None # 向量路名次（1 起） + vec_distance: float | None = None # 向量余弦距离，观测用 + kw_rank: int | None = None # 关键词路名次（1 起） + kw_tsrank: float | None = None # ts_rank 原始分，观测用 + + +@dataclass(frozen=True, slots=True) +class ScoredChunk: + """重排后的最终结果。各分量都归一到 [0, 1]，final_score 是加权和。""" + + chunk_id: int + entity_id: int + chunk_type: str + text: str + vec_score: float + kw_score: float + rule_score: float + freshness_score: float + final_score: float + + +def rrf_component(rank: int | None, k: int = RRF_K) -> float: + """单路 RRF 分量 ``1/(k+rank)``，再除以 ``1/(k+1)`` 归一到 (0, 1]. + + rank=1 → 1.0，rank=k+1 → 0.5。若直接用原始倒数，0.45×(1/61) ≈ 0.007 + 与规则项 0.25×1.0 完全不可比，架构 §8.3 的权重就失去"相对重要性"的 + 含义 —— 归一化只是线性缩放，不改变单路内的排序。 + """ + if rank is None or rank <= 0: + return 0.0 + return (k + 1) / (k + rank) + + +def rrf_score(vec_rank: int | None, kw_rank: int | None, k: int = RRF_K) -> float: + """经典 RRF 原始分：``1/(k+rank_vec) + 1/(k+rank_kw)``（观测/调试用）. + + 未被某路召回的 chunk 只累计被召回那一路的倒数。 + """ + total = 0.0 + for rank in (vec_rank, kw_rank): + if rank is not None and rank > 0: + total += 1.0 / (k + rank) + return total + + +def freshness_score(updated_at: datetime, now: datetime | None = None) -> float: + """新鲜度：``0.5 ** (age_days / 半衰期)``，越新越接近 1. + + ``now`` 缺省取当前 UTC 时刻；显式传入是为了让单测能构造已知年龄。 + 未来时间戳（时钟漂移）按 0 天处理，得 1.0。 + """ + current = now if now is not None else datetime.now(UTC) + age_seconds = (current - updated_at).total_seconds() + age_days = max(age_seconds, 0.0) / 86400.0 + # float ** float 在 typeshed 里是 Any（负底数分数幂会出 complex），显式收窄。 + return float(0.5 ** (age_days / FRESHNESS_HALF_LIFE_DAYS)) + + +def rerank( + rows: Sequence[RecallRow], + *, + weights: Mapping[str, float] | None = None, + rule_boost: Mapping[int, float] | None = None, + now: datetime | None = None, + top_k: int | None = None, +) -> list[ScoredChunk]: + """融合重排：``final = w_v×vec + w_k×kw + w_r×rule + w_f×fresh``（架构 §8.3）. + + - ``vec``/``kw`` 是归一化的 RRF 分量（见 :func:`rrf_component`）； + - ``rule`` 直接取 ``rule_boost[chunk_id]``（调用方传入的规则权重， + 规则引擎联动属后续任务，管线只认接口）； + - ``fresh`` 是半衰期衰减分。 + + 排序稳定：final_score 降序，同分按 chunk_id 升序 —— 同一数据集重跑 + 结果可复现。``top_k`` 为 ``None`` 时不截断。 + """ + resolved_weights = weights if weights is not None else DEFAULT_RERANK_WEIGHTS + w_vec = float(resolved_weights.get("vector", 0.0)) + w_kw = float(resolved_weights.get("keyword", 0.0)) + w_rule = float(resolved_weights.get("rule", 0.0)) + w_fresh = float(resolved_weights.get("freshness", 0.0)) + boosts = rule_boost or {} + + out = [] + for row in rows: + vec = rrf_component(row.vec_rank) + kw = rrf_component(row.kw_rank) + rule = float(boosts.get(row.chunk_id, 0.0)) + fresh = freshness_score(row.updated_at, now) + out.append( + ScoredChunk( + chunk_id=row.chunk_id, + entity_id=row.entity_id, + chunk_type=row.chunk_type, + text=row.text, + vec_score=vec, + kw_score=kw, + rule_score=rule, + freshness_score=fresh, + final_score=w_vec * vec + w_kw * kw + w_rule * rule + w_fresh * fresh, + ) + ) + out.sort(key=lambda c: (-c.final_score, c.chunk_id)) + return out[:top_k] if top_k is not None else out diff --git a/packages/retrieval/service.py b/packages/retrieval/service.py new file mode 100644 index 0000000..b561a07 --- /dev/null +++ b/packages/retrieval/service.py @@ -0,0 +1,114 @@ +"""检索编排：embed → 混合检索 →（无结果时）缺口埋点（架构 §8.3）. + +只做编排不做策略：模型、权重、top_k 全部来自 Settings；``rule_boost`` 的 +生产者（规则引擎联动）属后续任务，这里只认接口。检索是内部能力， +没有 HTTP 路由 —— 被 agent / report 的 service 调用。 +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from packages.retrieval import repo +from packages.retrieval.embedding import EmbeddingProvider, build_embedding_provider +from packages.retrieval.rerank import DEFAULT_RERANK_WEIGHTS, ScoredChunk +from packages.shared.config import Settings +from packages.shared.logging import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen=True, slots=True) +class UserContext: + """调用方给的检索上下文（用户画像/场景），驱动硬过滤与规则加权。 + + ``body_shape`` 是禁忌过滤的输入：命中 ``avoid_body_shapes`` 的实体 + 绝不出现（FR-KB-06），而不是排得靠后。 + """ + + body_shape: str | None = None + occasion: str | None = None + entity_types: tuple[str, ...] = () + rule_boost: Mapping[int, float] = field(default_factory=dict) + + def filter_ctx(self) -> dict[str, Any]: + """落进 kb_gap 的过滤条件快照：缺口是在什么条件下出现的。""" + return { + "body_shape": self.body_shape, + "occasion": self.occasion, + "entity_types": list(self.entity_types), + } + + +@dataclass(frozen=True, slots=True) +class SearchOutcome: + """一次检索的结果。``degraded`` 表示无结果并已记入 kb_gap（T-P2-06）。""" + + results: list[ScoredChunk] + degraded: bool + + +class RetrievalService: + """混合检索的编排入口（嵌入 + 检索 + 缺口埋点）.""" + + def __init__( + self, + session: AsyncSession, + *, + embedding: EmbeddingProvider, + weights: Mapping[str, float] | None = None, + top_k: int = 50, + ) -> None: + self._session = session + self._embedding = embedding + self._weights: dict[str, float] = ( + dict(weights) if weights is not None else dict(DEFAULT_RERANK_WEIGHTS) + ) + self._top_k = top_k + + @classmethod + def from_settings(cls, session: AsyncSession, settings: Settings) -> RetrievalService: + """从 Settings 装配（模型、权重、top_k 一处配置，换模型不发版）.""" + return cls( + session, + embedding=build_embedding_provider(settings), + weights=settings.retrieval_rerank_weights, + top_k=settings.retrieval_top_k, + ) + + async def search( + self, query_text: str, *, user_ctx: UserContext | None = None + ) -> SearchOutcome: + ctx = user_ctx or UserContext() + # embed 的输入就是 query 原文（DeterministicEmbedder/未来的 BGE-M3 都 + # 不需要拼提示词），保持「同 query 同向量」，缓存与复跑才有意义。 + query_vec = (await self._embedding.embed([query_text]))[0] + results = await repo.hybrid_search( + self._session, + query_vec, + query_text, + entity_types=ctx.entity_types or None, + occasion=ctx.occasion, + body_shape=ctx.body_shape, + rule_boost=ctx.rule_boost, + weights=self._weights, + top_k=self._top_k, + ) + if results: + return SearchOutcome(results=results, degraded=False) + + # 无结果也是信号：落 kb_gap 驱动内容生产（T-P2-06）。埋点失败不该 + # 让调用方拿到异常 —— 检索空结果本身是合法结局，这里吞掉埋点异常 + # 只记日志。query 原文进表不进日志（日志里只留长度与过滤条件）。 + try: + await repo.record_gap(self._session, query_text, ctx.filter_ctx()) + await self._session.commit() + except Exception: + await self._session.rollback() + logger.error("retrieval_gap_record_failed", query_len=len(query_text)) + logger.info("retrieval_gap", query_len=len(query_text), **ctx.filter_ctx()) + return SearchOutcome(results=[], degraded=True) diff --git a/tests/integration/test_hybrid_retrieval.py b/tests/integration/test_hybrid_retrieval.py new file mode 100644 index 0000000..9919f83 --- /dev/null +++ b/tests/integration/test_hybrid_retrieval.py @@ -0,0 +1,297 @@ +"""混合检索集成：真实 PostgreSQL + pgvector，嵌入用 DeterministicEmbedder（不联网）. + +证明三件单测证明不了的事（任务文档 §4）： +1. 硬过滤零漏 —— 禁忌/draft 实体即使向量最近也绝不出现（T-P2-04 红线）； +2. 关键词路径在 tsv 'simple' 分词下真实命中； +3. 无结果 query 落 kb_gap 且 hit_count 递增（T-P2-06）。 +外加：向量化重跑整组替换且逐行带版本（T-P2-03）、kb_asset partial 索引可用。 + +Run: MEIDA_RUN_INTEGRATION=1 .venv/bin/pytest -m integration -q --no-cov +""" + +from __future__ import annotations + +import os +import uuid +from collections.abc import AsyncIterator +from typing import Any + +import pytest +from sqlalchemy import delete, select, text + +from packages.domain.ids import next_id +from packages.retrieval.embedding import DeterministicEmbedder +from packages.retrieval.repo import ChunkInput, hybrid_search, upsert_chunks +from packages.retrieval.service import RetrievalService, UserContext +from packages.shared.db import dispose_engine, session_scope +from packages.shared.models import KbAsset, KbChunk, KbEntity, KbGap + + +def _infra_available() -> bool: + return bool(os.getenv("MEIDA_RUN_INTEGRATION")) + + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + not _infra_available(), + reason="set MEIDA_RUN_INTEGRATION=1 with PostgreSQL and Redis running", + ), +] + + +def _unique(name: str) -> str: + return f"it-t8-{name}-{uuid.uuid4().hex[:10]}" + + +@pytest.fixture(autouse=True) +async def _hermetic_engine() -> AsyncIterator[None]: + """每个用例前后重建全局 engine：asyncpg 连接绑定创建它的事件循环， + pytest-asyncio 给每个用例一个新 loop，复用池里旧连接会报 + "attached to a different loop"。app 系用例由 lifespan 的 dispose_engine + 兜底；直接用 session_scope 的本文件必须自己管。""" + await dispose_engine() + yield + await dispose_engine() + + +@pytest.fixture +async def clean_kb() -> AsyncIterator[dict[str, list[Any]]]: + """登记本用例造的脏数据并在结束时清掉：集成库长活，不能留假实体.""" + registry: dict[str, list[Any]] = {"entities": [], "gaps": []} + yield registry + async with session_scope() as session: + if registry["entities"]: + ids = registry["entities"] + await session.execute(delete(KbChunk).where(KbChunk.entity_id.in_(ids))) + await session.execute(delete(KbAsset).where(KbAsset.entity_id.in_(ids))) + await session.execute(delete(KbEntity).where(KbEntity.entity_id.in_(ids))) + for query_text in registry["gaps"]: + await session.execute(delete(KbGap).where(KbGap.query_text == query_text)) + + +async def _add_entity( + registry: dict[str, list[Any]], + *, + entity_type: str = "item", + status: str = "published", + attrs: dict[str, Any] | None = None, +) -> int: + entity_id = next_id() + registry["entities"].append(entity_id) + async with session_scope() as session: + session.add( + KbEntity( + entity_id=entity_id, + entity_type=entity_type, + name=_unique(entity_type), + attrs=attrs or {}, + status=status, + ) + ) + return entity_id + + +class RenamedEmbedder: + """换名复用确定性向量：验证 model/version 逐行落库与整组替换（T-P2-03）.""" + + def __init__(self, model_name: str, version: str) -> None: + self._inner = DeterministicEmbedder(dimensions=1024) + self.model_name = model_name + self.version = version + self.dimensions = 1024 + + async def embed(self, texts: list[str]) -> list[list[float]]: + return await self._inner.embed(texts) + + +# --- 1. 硬过滤零漏（T-P2-04 红线）-------------------------------------------- +async def test_hard_filter_never_leaks_taboo_or_draft(clean_kb: dict[str, list[Any]]) -> None: + embedder = DeterministicEmbedder(dimensions=1024) + query_text = "pear shape body dress advice" + # 防御性登记：若回归导致空结果被记进 kb_gap，用例结束也能清干净。 + clean_kb["gaps"].append(query_text) + + taboo_id = await _add_entity( + clean_kb, attrs={"occasions": ["wedding"], "avoid_body_shapes": ["apple"]} + ) + draft_id = await _add_entity(clean_kb, status="draft") + good_id = await _add_entity(clean_kb, attrs={"occasions": ["wedding"]}) + + async with session_scope() as session: + # 禁忌与 draft 的 chunk 和 query 完全同文：不过滤的话它们必然霸占 + # 向量第一名（距离 0）—— 这样断言"零出现"才有说服力。 + await upsert_chunks(session, taboo_id, [ChunkInput(text=query_text)], embedder) + await upsert_chunks(session, draft_id, [ChunkInput(text=query_text)], embedder) + await upsert_chunks( + session, + good_id, + [ChunkInput(text="soft knit sweater paired with a high waist skirt")], + embedder, + ) + + async with session_scope() as session: + svc = RetrievalService(session, embedding=embedder) + outcome = await svc.search( + query_text, + user_ctx=UserContext(body_shape="apple", occasion="wedding", entity_types=("item",)), + ) + + assert outcome.degraded is False + assert outcome.results, "正常 published 实体应被召回" + assert {r.entity_id for r in outcome.results} == {good_id} + + # 对照组：去掉禁忌/场景过滤后禁忌实体确实在库里且可被召回 —— + # 证明上面拦截它的是硬过滤，而不是数据没写进去。draft 永远不出现。 + query_vec = (await embedder.embed([query_text]))[0] + async with session_scope() as session: + control = await hybrid_search(session, query_vec, query_text) + control_entities = {r.entity_id for r in control} + assert taboo_id in control_entities + assert draft_id not in control_entities + + +# --- 2. 关键词路径（tsv 'simple' 分词）---------------------------------------- +async def test_keyword_recall_ranks_matching_chunk_first( + clean_kb: dict[str, list[Any]], +) -> None: + embedder = DeterministicEmbedder(dimensions=1024) + entity_id = await _add_entity(clean_kb) + matching = "breathable linen shirt keeps you cool in summer" + unrelated = "heavy wool coat for freezing winter nights" + + async with session_scope() as session: + await upsert_chunks( + session, + entity_id, + [ChunkInput(text=matching), ChunkInput(text=unrelated)], + embedder, + ) + + # 查询文本与两个 chunk 都不同文 → 向量路分不出高下（近似正交）， + # 排序由关键词路（0.2 权重、归一后差距 1.0）决定。 + async with session_scope() as session: + results = await hybrid_search( + session, (await embedder.embed(["linen shirt summer"]))[0], "linen shirt summer" + ) + + assert [r.text for r in results] == [matching, unrelated] + assert results[0].kw_score == pytest.approx(1.0) + assert results[1].kw_score == 0.0 + + +# --- 3. 向量化可重跑（T-P2-03）------------------------------------------------ +async def test_upsert_chunks_is_versioned_and_rerunnable( + clean_kb: dict[str, list[Any]], +) -> None: + entity_id = await _add_entity(clean_kb) + + async with session_scope() as session: + first = await upsert_chunks( + session, + entity_id, + [ChunkInput(text="rule one"), ChunkInput(text="rule two")], + RenamedEmbedder("bge-m3", "v1"), + ) + assert {c.embedding_model for c in first} == {"bge-m3"} + assert {c.embedding_version for c in first} == {"v1"} + + # 换模型名（= 换版本）重跑：整组替换，不残留旧版本向量 + async with session_scope() as session: + await upsert_chunks( + session, + entity_id, + [ChunkInput(text="rule one"), ChunkInput(text="rule two"), ChunkInput(text="rule new")], + RenamedEmbedder("bge-m4", "v2"), + ) + + async with session_scope() as session: + rows = ( + (await session.execute(select(KbChunk).where(KbChunk.entity_id == entity_id))) + .scalars() + .all() + ) + assert len(rows) == 3 + assert {r.embedding_model for r in rows} == {"bge-m4"} + assert {r.embedding_version for r in rows} == {"v2"} + + +# --- 4. 无结果 query 落 kb_gap（T-P2-06）-------------------------------------- +async def test_no_result_query_lands_in_kb_gap_with_counter( + clean_kb: dict[str, list[Any]], +) -> None: + query_text = _unique("missing") + clean_kb["gaps"].append(query_text) + svc_embedder = DeterministicEmbedder(dimensions=1024) + + async with session_scope() as session: + svc = RetrievalService(session, embedding=svc_embedder) + first = await svc.search(query_text, user_ctx=UserContext(entity_types=("nonexistent",))) + second = await svc.search(query_text, user_ctx=UserContext(entity_types=("nonexistent",))) + + assert first.degraded and first.results == [] + assert second.degraded and second.results == [] + + async with session_scope() as session: + gaps = ( + (await session.execute(select(KbGap).where(KbGap.query_text == query_text))) + .scalars() + .all() + ) + assert len(gaps) == 1 # 同 query + 同 filter_ctx 只有一行 + assert gaps[0].hit_count == 2 + assert gaps[0].filter_ctx["entity_types"] == ["nonexistent"] + + +# --- 5. kb_asset 按 entity 查询 + partial 索引可用 ----------------------------- +async def test_tryon_ready_assets_query_uses_partial_index( + clean_kb: dict[str, list[Any]], +) -> None: + entity_id = await _add_entity(clean_kb) + async with session_scope() as session: + session.add_all( + [ + KbAsset( + asset_id=next_id(), + entity_id=entity_id, + asset_type="front", + cos_key=_unique("front"), + tryon_ready=True, + review_status="approved", + ), + KbAsset( + asset_id=next_id(), + entity_id=entity_id, + asset_type="side", + cos_key=_unique("side"), + tryon_ready=False, + review_status="pending", + ), + ] + ) + + async with session_scope() as session: + ready = ( + ( + await session.execute( + select(KbAsset).where( + KbAsset.entity_id == entity_id, KbAsset.tryon_ready.is_(True) + ) + ) + ) + .scalars() + .all() + ) + assert len(ready) == 1 + assert ready[0].tryon_ready is True + + # 关闭顺序扫描逼 planner 走索引，证明 partial 索引真的可用 + # （SET LOCAL 只影响当前事务，随 session_scope 提交即失效）。 + await session.execute(text("SET LOCAL enable_seqscan = off")) + plan_lines = ( + (await session.execute(text("EXPLAIN SELECT asset_id FROM kb_asset WHERE tryon_ready"))) + .scalars() + .all() + ) + plan_text = "\n".join(str(line) for line in plan_lines) + assert "ix_kb_asset_tryon_ready" in plan_text diff --git a/tests/unit/test_embedding.py b/tests/unit/test_embedding.py new file mode 100644 index 0000000..ae055de --- /dev/null +++ b/tests/unit/test_embedding.py @@ -0,0 +1,99 @@ +"""DeterministicEmbedder 与工厂：纯数学，不连库不联网（T-P2-03 接口位）.""" + +from __future__ import annotations + +import math + +import pytest + +from packages.retrieval.embedding import ( + DETERMINISTIC_MODEL_NAME, + DeterministicEmbedder, + build_embedding_provider, +) +from packages.shared.config import get_settings +from packages.shared.errors import AppError, ErrorCode + + +def _dot(a: list[float], b: list[float]) -> float: + return sum(x * y for x, y in zip(a, b, strict=True)) + + +async def test_same_text_embeds_to_same_vector() -> None: + embedder = DeterministicEmbedder(dimensions=1024) + first, second = await embedder.embed(["白衬衫怎么搭配", "白衬衫怎么搭配"]) + assert first == second + + +async def test_dimensions_match_declaration() -> None: + vec = (await DeterministicEmbedder(dimensions=1024).embed(["hello"]))[0] + assert len(vec) == 1024 + small = (await DeterministicEmbedder(dimensions=8).embed(["hello"]))[0] + assert len(small) == 8 + + +async def test_vectors_are_l2_normalized() -> None: + for text in ("hello", "中文知识片段", ""): + vec = (await DeterministicEmbedder(dimensions=1024).embed([text]))[0] + assert math.isclose(_dot(vec, vec), 1.0, rel_tol=1e-6, abs_tol=1e-9) + + +async def test_different_texts_are_near_orthogonal() -> None: + """不同文本近似正交（|dot| ≪ 1）：向量召回管线在无真模型时仍可测.""" + + left, right = await DeterministicEmbedder(dimensions=1024).embed( + ["interview outfit advice", "winter skincare routine"] + ) + assert left != right + assert abs(_dot(left, right)) < 0.1 + + +async def test_batch_preserves_order_and_length() -> None: + texts = ["alpha", "beta", "gamma"] + vectors = await DeterministicEmbedder(dimensions=16).embed(texts) + assert len(vectors) == len(texts) + singles = [(await DeterministicEmbedder(dimensions=16).embed([t]))[0] for t in texts] + assert vectors == singles + + +def test_metadata_is_stable() -> None: + embedder = DeterministicEmbedder(dimensions=8) + assert embedder.model_name == DETERMINISTIC_MODEL_NAME + assert embedder.version + assert embedder.dimensions == 8 + + +def test_invalid_dimensions_rejected() -> None: + with pytest.raises(ValueError, match="dimensions"): + DeterministicEmbedder(dimensions=0) + + +def test_factory_builds_deterministic_embedder(monkeypatch: pytest.MonkeyPatch) -> None: + # 带 alias 的 Settings 字段按字段名构造会被静默忽略 —— 统一走 + # setenv + cache_clear（上批经验，见任务文档 §2）。 + monkeypatch.setenv("KB_EMBEDDING_MODEL", "deterministic-v1") + monkeypatch.setenv("KB_VECTOR_DIMENSIONS", "1024") + get_settings.cache_clear() + try: + settings = get_settings() + provider = build_embedding_provider(settings) + assert isinstance(provider, DeterministicEmbedder) + assert provider.model_name == "deterministic-v1" + assert provider.dimensions == 1024 + finally: + monkeypatch.undo() + get_settings.cache_clear() + + +def test_factory_rejects_unknown_model(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("KB_EMBEDDING_MODEL", "bge-m3-onnx") + get_settings.cache_clear() + try: + settings = get_settings() + assert settings.kb_embedding_model == "bge-m3-onnx" + with pytest.raises(AppError) as excinfo: + build_embedding_provider(settings) + assert excinfo.value.code == ErrorCode.CONFIG_ERROR + finally: + monkeypatch.undo() + get_settings.cache_clear() diff --git a/tests/unit/test_rerank.py b/tests/unit/test_rerank.py new file mode 100644 index 0000000..cabd715 --- /dev/null +++ b/tests/unit/test_rerank.py @@ -0,0 +1,158 @@ +"""RRF 融合 + 权重 + 新鲜度的数学正确性（T-P2-04 验收 3）：构造已知名次断言最终序.""" + +from __future__ import annotations + +import itertools +from datetime import UTC, datetime, timedelta + +import pytest + +from packages.retrieval.rerank import ( + DEFAULT_RERANK_WEIGHTS, + FRESHNESS_HALF_LIFE_DAYS, + RRF_K, + RecallRow, + freshness_score, + rerank, + rrf_component, + rrf_score, +) + +NOW = datetime(2026, 8, 14, 12, 0, 0, tzinfo=UTC) + + +def _row( + chunk_id: int, + *, + vec_rank: int | None = None, + kw_rank: int | None = None, + updated_at: datetime = NOW, + entity_id: int = 1, +) -> RecallRow: + return RecallRow( + chunk_id=chunk_id, + entity_id=entity_id, + chunk_type="text", + text=f"chunk-{chunk_id}", + updated_at=updated_at, + vec_rank=vec_rank, + kw_rank=kw_rank, + ) + + +# --- 单路 RRF 分量 ------------------------------------------------------------ +def test_rrf_component_is_max_normalized() -> None: + assert rrf_component(1) == pytest.approx(1.0) # (60+1)/(60+1) + assert rrf_component(2) == pytest.approx(61 / 62) + assert rrf_component(RRF_K + 2) == pytest.approx(0.5) # 61/122：名次再排后一半 + assert rrf_component(None) == 0.0 + assert rrf_component(0) == 0.0 + assert rrf_component(-3) == 0.0 + + +def test_rrf_component_monotonic_non_increasing() -> None: + values = [rrf_component(r) for r in range(1, 20)] + assert all(a >= b for a, b in itertools.pairwise(values)) + + +def test_rrf_score_is_sum_of_raw_reciprocals() -> None: + assert rrf_score(1, 2) == pytest.approx(1 / 61 + 1 / 62) + assert rrf_score(None, 1) == pytest.approx(1 / 61) + assert rrf_score(1, None) == pytest.approx(1 / 61) + assert rrf_score(None, None) == 0.0 + + +# --- 新鲜度 ------------------------------------------------------------------ +def test_freshness_half_life_decay() -> None: + assert freshness_score(NOW, now=NOW) == pytest.approx(1.0) + old = NOW - timedelta(days=FRESHNESS_HALF_LIFE_DAYS) + assert freshness_score(old, now=NOW) == pytest.approx(0.5) + older = NOW - timedelta(days=2 * FRESHNESS_HALF_LIFE_DAYS) + assert freshness_score(older, now=NOW) == pytest.approx(0.25) + # 时钟漂移产生未来时间戳 → 按 0 天处理，不放大分数 + future = NOW + timedelta(hours=1) + assert freshness_score(future, now=NOW) == pytest.approx(1.0) + + +# --- 加权公式 ---------------------------------------------------------------- +def test_final_score_follows_architecture_formula() -> None: + """final = 0.45×vec + 0.2×kw + 0.25×rule + 0.1×fresh，逐项核对.""" + + row = _row(1, vec_rank=1, kw_rank=3) + boost = {1: 0.5} + scored = rerank([row], rule_boost=boost, now=NOW) + expected = 0.45 * 1.0 + 0.2 * (61 / 63) + 0.25 * 0.5 + 0.1 * 1.0 + assert scored[0].final_score == pytest.approx(expected) + assert scored[0].vec_score == pytest.approx(1.0) + assert scored[0].kw_score == pytest.approx(61 / 63) + assert scored[0].rule_score == pytest.approx(0.5) + assert scored[0].freshness_score == pytest.approx(1.0) + + +def test_default_weights_match_architecture_section_8_3() -> None: + assert DEFAULT_RERANK_WEIGHTS == { + "vector": 0.45, + "keyword": 0.2, + "rule": 0.25, + "freshness": 0.1, + } + + +def test_dual_recall_beats_single_recall() -> None: + """两路都命中的 chunk 应压过只有向量路命中、名次更高的 chunk.""" + + both = _row(1, vec_rank=5, kw_rank=2) + vec_only = _row(2, vec_rank=1) + scored = rerank([both, vec_only], now=NOW) + by_id = {c.chunk_id: c for c in scored} + assert by_id[1].final_score > by_id[2].final_score + + +def test_rule_boost_lifts_lower_recall_chunk() -> None: + """规则权重 0.25×boost 可以把召回名次落后的 chunk 抬到最前（自建知识库的意义）.""" + + rank5 = _row(1, vec_rank=5) + rank1 = _row(2, vec_rank=1) + scored = rerank([rank5, rank1], rule_boost={1: 1.0}, now=NOW) + assert scored[0].chunk_id == 1 + assert scored[0].rule_score == pytest.approx(1.0) + + +def test_keyword_only_chunk_gets_zero_vector_score() -> None: + scored = rerank([_row(7, kw_rank=1)], now=NOW) + assert scored[0].vec_score == 0.0 + assert scored[0].kw_score == pytest.approx(1.0) + + +def test_freshness_breaks_ties_toward_newer_chunk() -> None: + """名次相同（现实中少见，但公式上可能）时，更新鲜者胜.""" + + fresh = _row(1, vec_rank=2, updated_at=NOW) + stale = _row(2, vec_rank=2, updated_at=NOW - timedelta(days=90)) + scored = rerank([fresh, stale], now=NOW) + assert scored[0].chunk_id == 1 + delta = scored[0].final_score - scored[1].final_score + assert delta == pytest.approx(0.1 * (1.0 - 0.5**3)) + + +def test_zero_weight_dimension_ignores_that_signal() -> None: + old = _row(1, vec_rank=1, updated_at=NOW - timedelta(days=365)) + new = _row(2, vec_rank=2, updated_at=NOW) + scored = rerank([old, new], weights={"vector": 1.0}, now=NOW) + assert scored[0].chunk_id == 1 # 新鲜度权重为 0，旧的反而名次靠前 + + +def test_stable_ordering_and_top_k_truncation() -> None: + rows = [_row(i, vec_rank=i) for i in (30, 10, 20, 40)] + scored = rerank(rows, now=NOW, top_k=2) + assert [c.chunk_id for c in scored] == [10, 20] + + +def test_equal_final_scores_break_tie_by_chunk_id() -> None: + rows = [_row(90, kw_rank=1), _row(10, kw_rank=1)] + scored = rerank(rows, now=NOW) + assert [c.chunk_id for c in scored] == [10, 90] + + +def test_empty_input_returns_empty() -> None: + assert rerank([], now=NOW) == [] diff --git a/tests/unit/test_retrieval_service.py b/tests/unit/test_retrieval_service.py new file mode 100644 index 0000000..4a6ab72 --- /dev/null +++ b/tests/unit/test_retrieval_service.py @@ -0,0 +1,179 @@ +"""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