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 完全不可比）。
159 lines
5.4 KiB
Python
159 lines
5.4 KiB
Python
"""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) == []
|