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)
|
||||
Reference in New Issue
Block a user