xsl
9c1802130d
第三批（T12 后台/T13 黄金集/T14 检索缓存）协调提交： - kb：KbAssetService（登记/审核，tryon_ready 只能随 approved 打开）+ Asset 模型与 repo - retrieval：list_gaps 缺口看板查询 + invalidate_cache() 钩子（T14 背后实现真缓存） - 配置：ADMIN_TOKEN（空=后台 404）+ 检索缓存三参数 - pytest golden marker（默认跳过，CI 强制跑）+ Makefile golden + CI 防线 5c - 测试隔离修复：test_auth_sms 加 autouse 清 sms:* 键夹具（频控计数跨运行 累积会毒化后续运行——随机 IP 空间仅 254 且部分用例固定 device_id）
285 lines
11 KiB
Python
285 lines
11 KiB
Python
"""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)
|
||
|
||
|
||
#: 进程内缓存代数：invalidate_cache() 自增，缓存键带上它即可整体失效（T-P2-10）。
|
||
#: 放 repo 而非 service，是因为 T14 的缓存模块会同文件引用，避免循环 import。
|
||
INVALIDATION_VERSION = 1
|
||
|
||
|
||
async def list_gaps(session: AsyncSession, *, limit: int = 50, offset: int = 0) -> list[RowMapping]:
|
||
"""缺口看板查询：hit_count 降序，同频次按最近出现降序。"""
|
||
from sqlalchemy import select
|
||
|
||
from packages.shared.models import KbGap
|
||
|
||
stmt = (
|
||
select(KbGap)
|
||
.order_by(KbGap.hit_count.desc(), KbGap.last_seen_at.desc())
|
||
.limit(limit)
|
||
.offset(offset)
|
||
)
|
||
result = await session.execute(stmt)
|
||
return list(result.mappings().all())
|