feat(retrieval): hybrid search over pgvector with hard filters and gap tracking
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 完全不可比）。
This commit is contained in:
@@ -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）。
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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},
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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) == []
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user