meida_backend/packages/retrieval/service.py

"""检索编排：embed → 混合检索 →（无结果时）缺口埋点（架构 §8.3）.

只做编排不做策略：模型、权重、top_k 全部来自 Settings；``rule_boost`` 的
生产者（规则引擎联动）属后续任务，这里只认接口。检索是内部能力，
没有 HTTP 路由 —— 被 agent / report 的 service 调用。

T-P2-10：有结果的检索进进程内缓存（TTL+LRU，见 cache.py），空结果
不缓存（gap 埋点每次都要跑）；发布/归档后调 :func:`invalidate_cache`
整体失效。
"""

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.cache import CacheKey, TtlCache
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


# --- 进程级共享缓存（T-P2-10）-------------------------------------------------
#: RetrievalService 每请求新建（各自绑定 session），但缓存的价值在**跨请求**
#: 复用 —— 实例必须模块级共享。惰性创建；Settings 的 TTL/容量变化时重建
#: （进程内配置恒定，重建只发生在测试或热改配置的场景）。
_shared_cache: TtlCache[SearchOutcome] | None = None


def _get_shared_cache(settings: Settings) -> TtlCache[SearchOutcome] | None:
    """按 Settings 返回共享缓存；``retrieval_cache_enabled=False`` → None（禁用）."""
    global _shared_cache
    if not settings.retrieval_cache_enabled:
        return None
    if (
        _shared_cache is None
        or _shared_cache.ttl_seconds != settings.retrieval_cache_ttl_seconds
        or _shared_cache.max_entries != settings.retrieval_cache_max_entries
    ):
        _shared_cache = TtlCache(
            ttl_seconds=settings.retrieval_cache_ttl_seconds,
            max_entries=settings.retrieval_cache_max_entries,
        )
    return _shared_cache


class RetrievalService:
    """混合检索的编排入口（嵌入 + 检索 + 缺口埋点）."""

    def __init__(
        self,
        session: AsyncSession,
        *,
        embedding: EmbeddingProvider,
        weights: Mapping[str, float] | None = None,
        top_k: int = 50,
        cache: TtlCache[SearchOutcome] | None = None,
    ) -> 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
        # None = 禁用缓存，行为与没有缓存时完全一致（默认值，既有调用点不受影响）。
        self._cache = cache

    @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,
            cache=_get_shared_cache(settings),
        )

    async def search(
        self, query_text: str, *, user_ctx: UserContext | None = None
    ) -> SearchOutcome:
        ctx = user_ctx or UserContext()
        cache, key = self._cache, None
        if cache is not None:
            # 全输入进键（含 rule_boost 哈希与失效代数）——缺一个分量就会
            # 把 A 条件的结果串给 B 条件（cache.py 红线 1）。命中时 embed
            # 与 DB 全部跳过：缓存命中 = DB 零交互。
            key = CacheKey.build(
                version=repo.INVALIDATION_VERSION,
                query_text=query_text,
                body_shape=ctx.body_shape,
                occasion=ctx.occasion,
                entity_types=ctx.entity_types,
                rule_boost=ctx.rule_boost,
                top_k=self._top_k,
            )
            cached = cache.get(key)
            if cached is not None:
                # results 是 list（可变）：返回拷贝，调用方原地改 list 不可能
                # 污染缓存（cache.py 红线 2）。
                return SearchOutcome(results=list(cached.results), degraded=cached.degraded)
        # 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:
            if cache is not None and key is not None:
                # 只缓存有结果的检索。空结果不缓存：degraded 路径的 record_gap
                # 每次都要跑，缓存空结果会把缺口计数吞掉。存的是拷贝——返回给
                # 调用方的 list 与缓存里的不是同一个对象。
                cache.put(key, SearchOutcome(results=list(results), degraded=False))
            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)


class GapSummary:
    """kb_gap 的一行摘要（缺口看板的数据形状）。V1 用普通类：字段还会随
    看板需求演进，先不冻结成 dataclass。"""

    def __init__(self, row: Any) -> None:
        self.gap_id = row.gap_id
        self.query_text = row.query_text
        self.filter_ctx = row.filter_ctx
        self.hit_count = row.hit_count
        self.first_seen_at = row.first_seen_at
        self.last_seen_at = row.last_seen_at


async def list_gaps(session: AsyncSession, *, limit: int = 50, offset: int = 0) -> list[GapSummary]:
    """缺口看板：按 hit_count 降序（最常被问到的缺口排最前，T-P2-06）。"""
    rows = await repo.list_gaps(session, limit=limit, offset=offset)
    return [GapSummary(r) for r in rows]


def invalidate_cache() -> None:
    """发布/归档知识后调用：让检索结果缓存整体失效（T-P2-10）。

    代数失效：``repo.INVALIDATION_VERSION`` 自增，之后所有查找都用新代
    构键，旧代条目（含检索途中才被写回的）永不命中——无需遍历删除。
    共享缓存实例同时清空，只为立刻释放内存，失效正确性只依赖代数。
    签名不变：admin 发布编排（T12）直接调用，不需要了解内部机制。
    """
    global _shared_cache
    repo.INVALIDATION_VERSION += 1
    if _shared_cache is not None:
        _shared_cache.clear()