meida_backend/packages/retrieval/repo.py

"""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)