xsl
aa1287d83d
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 完全不可比）。
115 lines
4.3 KiB
Python
115 lines
4.3 KiB
Python
"""检索编排：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)
|