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 完全不可比）。
298 lines
11 KiB
Python
298 lines
11 KiB
Python
"""混合检索集成：真实 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
|