feat(retrieval): TTL+LRU result cache with generation-based invalidation
TtlCache(全输入进键/惰性过期/容量逐出/命中零 DB)+ search 接线 (空结果不缓存,保住 gap 埋点)+ invalidate_cache 真实现(T-P2-10)。 实现: - packages/retrieval/cache.py(新建):CacheKey 全输入进键(query/body_shape/ occasion/entity_types 排序归一/rule_boost 排序序列化哈希/top_k/失效代数); TtlCache 惰性过期 + OrderedDict LRU(超限 popitem 最旧)+ hits/misses/ evictions 计数与 stats();clock 可注入(单测免 sleep,默认 monotonic); asyncio 单循环不加锁的前提写在模块 docstring;正确性红线(只缓存不可变 值、命中零 DB)同处。纯逻辑,不 import 任何项目内模块(值类型由 service 以 PEP 695 泛型注入,避免循环依赖)。 - packages/retrieval/service.py:__init__ 加 keyword-only cache=None(默认 禁用,行为与现状一致,既有调用点不受影响);from_settings 按 Settings 装配进程级共享缓存(retrieval_cache_enabled=False → None;TTL/容量变化 时重建);search() 在 embed 之前查缓存——命中直接返回(零 DB、零 embed), 写入与命中都做防御性拷贝(results 是 list,调用方拿不到缓存内引用), 仅有结果时写缓存(空结果不缓存,record_gap 每次都跑);invalidate_cache() 签名不变:自增 repo.INVALIDATION_VERSION(代数失效,旧代键永不命中, 与 T12 的契约保持)+ 清空共享实例(仅为立即回收内存,失效正确性只 依赖代数)。repo.py 未做任何改动。 测试: - 单测 tests/unit/test_retrieval_cache.py(21 例,fake repo/embedding + 注入时钟,不连库):同参二连查零 DB 且结果相等、键分量矩阵(query/ occasion/body_shape/entity_types/rule_boost/top_k 任一变化 → miss)、 TTL 过期边界(60.0 过期/59.9 命中)、LRU 逐出与命中率/计数、空结果 不缓存且 gap 埋点每次触发、invalidate_cache() 代数失效(含自增语义 保留断言)、cache=None 与现状一致、from_settings 装配开关与共享实例、 rule_boost 哈希确定性。全量单测 494 passed(基线 473 + 新增 21)。 - 集成 tests/integration/test_retrieval_cache.py(2 例,真实 PG + pgvector): 发布新实体 + invalidate_cache() 后同参检索立即可见(失效前缓存仍只吐 旧实体,证明缓存在工作——T-P2-10 验收语义);不失效、TTL=1s 到期 (asyncio.sleep 1.1s)后新实体自然可见(双路径验收)。集成全量 40 passed。 偏差说明:无(invalidate_cache() 在代数失效之外顺带清空共享缓存实例, 失效正确性只依赖代数,清空仅为立即释放内存)。
This commit is contained in:
@@ -0,0 +1,170 @@
|
||||
"""检索结果缓存的纯逻辑:全输入缓存键 + TTL/LRU 淘汰(T-P2-10).
|
||||
|
||||
刻意与 service 分离:键构造与淘汰策略必须能被单测逐项断言,混进
|
||||
service 就会把缓存语义测试拖进 DB 打桩。本模块不 import 任何项目内
|
||||
模块(``SearchOutcome`` 由 service 以泛型参数注入,避免 service ↔
|
||||
cache 循环依赖),``INVALIDATION_VERSION`` 只是键里的一个普通整数。
|
||||
|
||||
正确性红线(改本模块前先读):
|
||||
|
||||
1. **全部影响结果的输入都必须进缓存键**(见 :class:`CacheKey`)——缺
|
||||
一个分量就会把 A 条件的结果串给 B 条件。``rule_boost`` 是
|
||||
``dict[int, float]``,按键排序序列化再哈希;
|
||||
2. **只缓存不可变值**:值应是 frozen dataclass(``SearchOutcome``)。
|
||||
它内部的 ``results`` 是 list(可变)——写入与命中时 service 都做
|
||||
防御性拷贝,调用方永远拿不到缓存内 list 的引用,原地修改不可能
|
||||
污染缓存。绝不缓存 ORM 对象或任何可变引用;
|
||||
3. **命中路径零 DB**:缓存命中时不碰 session、不 embed、不重排,
|
||||
直接返回缓存值。
|
||||
|
||||
线程安全:**不加锁,也别加**。检索跑在 asyncio 单事件循环里,
|
||||
get/put 各自是一段无 await 的同步代码,事件循环只在 await 点切换
|
||||
协程,单循环内不存在竞态窗口。加锁反而暗示它可以跨线程使用——那
|
||||
是错误的使用方式;若未来把检索挪进线程池,必须重新评估本类。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
def rule_boost_hash(rule_boost: Mapping[int, float]) -> str:
|
||||
"""``rule_boost`` 的稳定哈希:按键排序序列化后 sha256.
|
||||
|
||||
dict 迭代顺序是插入序:``{1: 0.5, 2: 0.3}`` 与先放 2 再放 1 的同义
|
||||
dict 逐项迭代结果不同,必须排序后再序列化,否则同义输入得到不同
|
||||
键(白丢命中率)。float 用 ``repr`` 保留精确位,避免 ``str`` 截断
|
||||
把不同的 boost 折叠成同一个键。
|
||||
"""
|
||||
payload = ",".join(f"{cid!r}:{boost!r}" for cid, boost in sorted(rule_boost.items()))
|
||||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CacheKey:
|
||||
"""一次检索的全部输入指纹:任一分量不同就是另一次检索.
|
||||
|
||||
- ``version``:``repo.INVALIDATION_VERSION`` 的快照。发布/归档后
|
||||
自增,新查找全部用新代构键,旧代条目(含检索途中被失效后写回的)
|
||||
永不命中——这就是代数失效,不需要遍历删除;
|
||||
- ``entity_types`` 排序归一:SQL 的 ``= ANY(...)`` 与顺序无关,
|
||||
同集合不同顺序应共享同一条目;
|
||||
- ``rule_boost_hash``:见 :func:`rule_boost_hash`。
|
||||
"""
|
||||
|
||||
version: int
|
||||
query_text: str
|
||||
body_shape: str | None
|
||||
occasion: str | None
|
||||
entity_types: tuple[str, ...]
|
||||
rule_boost_hash: str
|
||||
top_k: int
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
*,
|
||||
version: int,
|
||||
query_text: str,
|
||||
top_k: int,
|
||||
body_shape: str | None = None,
|
||||
occasion: str | None = None,
|
||||
entity_types: Sequence[str] = (),
|
||||
rule_boost: Mapping[int, float] | None = None,
|
||||
) -> CacheKey:
|
||||
"""从 search 的入参构造键(``rule_boost`` 在这里收敛成哈希)."""
|
||||
return cls(
|
||||
version=version,
|
||||
query_text=query_text,
|
||||
body_shape=body_shape,
|
||||
occasion=occasion,
|
||||
entity_types=tuple(sorted(entity_types)),
|
||||
rule_boost_hash=rule_boost_hash(rule_boost or {}),
|
||||
top_k=top_k,
|
||||
)
|
||||
|
||||
|
||||
class TtlCache[V]:
|
||||
"""进程内 TTL + 容量 LRU 缓存:读时惰性过期,超限逐出最旧.
|
||||
|
||||
指标(``hits``/``misses``/``evictions``)先暴露成属性,Prometheus
|
||||
接入时直接读,:meth:`stats` 给出汇总快照。
|
||||
|
||||
线程安全前提见模块 docstring:asyncio 单循环使用,不加锁。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ttl_seconds: float,
|
||||
max_entries: int,
|
||||
clock: Callable[[], float] | None = None,
|
||||
) -> None:
|
||||
if ttl_seconds <= 0:
|
||||
raise ValueError(f"ttl_seconds 必须为正, 得到 {ttl_seconds}")
|
||||
if max_entries < 1:
|
||||
raise ValueError(f"max_entries 必须至少为 1, 得到 {max_entries}")
|
||||
self._ttl_seconds = ttl_seconds
|
||||
self._max_entries = max_entries
|
||||
# monotonic 而不是 wall clock:TTL 是时长语义,NTP 校时/回拨不该
|
||||
# 让缓存集体"突然过期"或"永不过期"。clock 可注入是为了单测免 sleep。
|
||||
self._clock = clock or time.monotonic
|
||||
self._entries: OrderedDict[CacheKey, tuple[float, V]] = OrderedDict()
|
||||
self.hits = 0
|
||||
self.misses = 0
|
||||
self.evictions = 0
|
||||
|
||||
@property
|
||||
def ttl_seconds(self) -> float:
|
||||
"""当前 TTL(秒)。from_settings 据此判断共享缓存是否需要重建."""
|
||||
return self._ttl_seconds
|
||||
|
||||
@property
|
||||
def max_entries(self) -> int:
|
||||
"""容量上限。用途同上."""
|
||||
return self._max_entries
|
||||
|
||||
def get(self, key: CacheKey) -> V | None:
|
||||
"""命中返回值并移到 LRU 末尾;未命中/已过期返回 ``None``.
|
||||
|
||||
过期是惰性的:读时发现超龄就地删除并按 miss 计——evictions 只
|
||||
描述容量压力(LRU 逐出),不掺入时间淘汰,两个数字语义分开。
|
||||
"""
|
||||
entry = self._entries.get(key)
|
||||
if entry is None:
|
||||
self.misses += 1
|
||||
return None
|
||||
stored_at, value = entry
|
||||
if self._clock() - stored_at >= self._ttl_seconds:
|
||||
del self._entries[key]
|
||||
self.misses += 1
|
||||
return None
|
||||
self._entries.move_to_end(key) # 刚被访问过,是最新的
|
||||
self.hits += 1
|
||||
return value
|
||||
|
||||
def put(self, key: CacheKey, value: V) -> None:
|
||||
"""写入/覆盖一个条目,并保证容量不超限(超出逐出最旧)."""
|
||||
self._entries[key] = (self._clock(), value)
|
||||
self._entries.move_to_end(key)
|
||||
while len(self._entries) > self._max_entries:
|
||||
self._entries.popitem(last=False)
|
||||
self.evictions += 1
|
||||
|
||||
def clear(self) -> None:
|
||||
"""整体清空。代数失效本身不需要它(旧代键永不命中),这里只为
|
||||
``invalidate_cache()`` 立刻释放内存,不等 TTL/LRU 兜底。"""
|
||||
self._entries.clear()
|
||||
|
||||
def stats(self) -> dict[str, int]:
|
||||
"""命中/未命中/逐出/当前条目数(观测快照)."""
|
||||
return {
|
||||
"hits": self.hits,
|
||||
"misses": self.misses,
|
||||
"evictions": self.evictions,
|
||||
"entries": len(self._entries),
|
||||
}
|
||||
@@ -3,6 +3,10 @@
|
||||
只做编排不做策略:模型、权重、top_k 全部来自 Settings;``rule_boost`` 的
|
||||
生产者(规则引擎联动)属后续任务,这里只认接口。检索是内部能力,
|
||||
没有 HTTP 路由 —— 被 agent / report 的 service 调用。
|
||||
|
||||
T-P2-10:有结果的检索进进程内缓存(TTL+LRU,见 cache.py),空结果
|
||||
不缓存(gap 埋点每次都要跑);发布/归档后调 :func:`invalidate_cache`
|
||||
整体失效。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -14,6 +18,7 @@ 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
|
||||
@@ -52,6 +57,30 @@ class SearchOutcome:
|
||||
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:
|
||||
"""混合检索的编排入口(嵌入 + 检索 + 缺口埋点)."""
|
||||
|
||||
@@ -62,6 +91,7 @@ class RetrievalService:
|
||||
embedding: EmbeddingProvider,
|
||||
weights: Mapping[str, float] | None = None,
|
||||
top_k: int = 50,
|
||||
cache: TtlCache[SearchOutcome] | None = None,
|
||||
) -> None:
|
||||
self._session = session
|
||||
self._embedding = embedding
|
||||
@@ -69,21 +99,43 @@ class RetrievalService:
|
||||
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 一处配置,换模型不发版)."""
|
||||
"""从 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]
|
||||
@@ -99,6 +151,11 @@ class RetrievalService:
|
||||
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)。埋点失败不该
|
||||
@@ -136,7 +193,12 @@ async def list_gaps(session: AsyncSession, *, limit: int = 50, offset: int = 0)
|
||||
def invalidate_cache() -> None:
|
||||
"""发布/归档知识后调用:让检索结果缓存整体失效(T-P2-10)。
|
||||
|
||||
第三批 T14 会在背后实现真正的进程内缓存;本桩保证调用方(运营后台的
|
||||
发布编排)现在就能接上这个钩子,T14 合并后无需再改调用点。
|
||||
代数失效:``repo.INVALIDATION_VERSION`` 自增,之后所有查找都用新代
|
||||
构键,旧代条目(含检索途中才被写回的)永不命中——无需遍历删除。
|
||||
共享缓存实例同时清空,只为立刻释放内存,失效正确性只依赖代数。
|
||||
签名不变:admin 发布编排(T12)直接调用,不需要了解内部机制。
|
||||
"""
|
||||
global _shared_cache
|
||||
repo.INVALIDATION_VERSION += 1
|
||||
if _shared_cache is not None:
|
||||
_shared_cache.clear()
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
"""检索缓存集成(T-P2-10):真实 PostgreSQL + pgvector,DeterministicEmbedder 不联网.
|
||||
|
||||
证明单测证明不了的两条失效路径(任务文档 §4,双路径验收):
|
||||
1. 发布新实体 + ``invalidate_cache()`` → 同参数检索立即可见新实体;
|
||||
2. 不失效、等 TTL 自然到期(TTL=1s,sleep 1.1s)→ 新实体同样可见。
|
||||
|
||||
Run: MEIDA_RUN_INTEGRATION=1 .venv/bin/pytest -m integration -q --no-cov
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import delete
|
||||
|
||||
from packages.domain.ids import next_id
|
||||
from packages.retrieval.cache import TtlCache
|
||||
from packages.retrieval.embedding import DeterministicEmbedder
|
||||
from packages.retrieval.repo import ChunkInput, upsert_chunks
|
||||
from packages.retrieval.service import RetrievalService, invalidate_cache
|
||||
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-t14-{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"。"""
|
||||
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 _publish_entity_with_chunk(
|
||||
registry: dict[str, list[Any]], embedder: DeterministicEmbedder, *, chunk_text: str
|
||||
) -> int:
|
||||
"""直接以 published 落库(任务文档允许「SQL 置 published」).
|
||||
|
||||
用例聚焦缓存失效语义,不走发布状态机;chunk 与查询同文保证向量路
|
||||
距离 0,必然被召回。
|
||||
"""
|
||||
entity_id = next_id()
|
||||
registry["entities"].append(entity_id)
|
||||
async with session_scope() as session:
|
||||
session.add(
|
||||
KbEntity(
|
||||
entity_id=entity_id,
|
||||
entity_type="item",
|
||||
name=_unique("item"),
|
||||
attrs={},
|
||||
status="published",
|
||||
)
|
||||
)
|
||||
# sessionmaker 是 autoflush=False:不显式 flush,挂起的实体 INSERT
|
||||
# 不会先于 kb_chunk 下发,FK 会拦下 chunk。
|
||||
await session.flush()
|
||||
await upsert_chunks(session, entity_id, [ChunkInput(text=chunk_text)], embedder)
|
||||
return entity_id
|
||||
|
||||
|
||||
def _seen_entities(outcome: Any) -> set[int]:
|
||||
return {r.entity_id for r in outcome.results}
|
||||
|
||||
|
||||
# --- 1. 发布 + invalidate_cache() → 新实体立即可见(T-P2-10 验收语义)---------
|
||||
async def test_publish_then_invalidate_makes_new_entity_visible(
|
||||
clean_kb: dict[str, list[Any]],
|
||||
) -> None:
|
||||
embedder = DeterministicEmbedder(dimensions=1024)
|
||||
query = _unique("invalidate-query")
|
||||
# 防御性登记:若回归导致空结果被记进 kb_gap,用例结束也能清干净。
|
||||
clean_kb["gaps"].append(query)
|
||||
|
||||
entity_a = await _publish_entity_with_chunk(clean_kb, embedder, chunk_text=query)
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=60.0, max_entries=64)
|
||||
|
||||
async with session_scope() as session:
|
||||
svc = RetrievalService(session, embedding=embedder, cache=cache)
|
||||
first = await svc.search(query)
|
||||
second = await svc.search(query)
|
||||
|
||||
assert first.degraded is False and second.degraded is False
|
||||
assert _seen_entities(first) == {entity_a}
|
||||
assert second == first
|
||||
assert cache.stats()["hits"] == 1 # 第二次命中缓存:零 DB、零 embed
|
||||
|
||||
entity_b = await _publish_entity_with_chunk(clean_kb, embedder, chunk_text=query)
|
||||
|
||||
# 未失效前缓存照旧只吐 A —— 证明缓存真的在工作(不是每次都查库)
|
||||
async with session_scope() as session:
|
||||
svc = RetrievalService(session, embedding=embedder, cache=cache)
|
||||
stale = await svc.search(query)
|
||||
assert _seen_entities(stale) == {entity_a}
|
||||
|
||||
invalidate_cache() # 发布编排(admin/T12)调用的钩子
|
||||
|
||||
async with session_scope() as session:
|
||||
svc = RetrievalService(session, embedding=embedder, cache=cache)
|
||||
fresh = await svc.search(query)
|
||||
assert _seen_entities(fresh) == {entity_a, entity_b} # 新实体立即可见
|
||||
|
||||
|
||||
# --- 2. 不失效,等 TTL 到期 → 新实体自然可见(双保险)-------------------------
|
||||
async def test_ttl_expiry_makes_new_entity_visible_without_invalidation(
|
||||
clean_kb: dict[str, list[Any]],
|
||||
) -> None:
|
||||
embedder = DeterministicEmbedder(dimensions=1024)
|
||||
query = _unique("ttl-query")
|
||||
clean_kb["gaps"].append(query)
|
||||
|
||||
entity_a = await _publish_entity_with_chunk(clean_kb, embedder, chunk_text=query)
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=1.0, max_entries=64)
|
||||
|
||||
async with session_scope() as session:
|
||||
svc = RetrievalService(session, embedding=embedder, cache=cache)
|
||||
first = await svc.search(query)
|
||||
await svc.search(query)
|
||||
assert _seen_entities(first) == {entity_a}
|
||||
assert cache.stats()["hits"] == 1 # TTL 内命中缓存
|
||||
|
||||
entity_b = await _publish_entity_with_chunk(clean_kb, embedder, chunk_text=query)
|
||||
# 不调用 invalidate_cache():只等 TTL 自然到期(1s + 余量 0.1s)。
|
||||
# asyncio.sleep 与 TtlCache 的默认时钟同为单调时钟,无校时跳变风险。
|
||||
await asyncio.sleep(1.1)
|
||||
|
||||
async with session_scope() as session:
|
||||
svc = RetrievalService(session, embedding=embedder, cache=cache)
|
||||
fresh = await svc.search(query)
|
||||
assert _seen_entities(fresh) == {entity_a, entity_b} # 到期后重新检索
|
||||
@@ -0,0 +1,359 @@
|
||||
"""检索结果缓存(TtlCache + search 接线):repo/embedding 全打桩,不连库.
|
||||
|
||||
覆盖任务文档 §4 的七组断言:命中零 DB、键分量矩阵、TTL 过期、LRU
|
||||
逐出、空结果不缓存(gap 埋点不被吞)、代数失效、cache=None 一致性;
|
||||
外加 rule_boost 哈希的确定性与 from_settings 的装配开关。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from packages.retrieval import repo
|
||||
from packages.retrieval import service as service_mod
|
||||
from packages.retrieval.cache import CacheKey, TtlCache, rule_boost_hash
|
||||
from packages.retrieval.rerank import ScoredChunk
|
||||
from packages.retrieval.service import RetrievalService, UserContext
|
||||
from packages.shared.config import Settings
|
||||
|
||||
|
||||
class FakeEmbedder:
|
||||
"""满足 EmbeddingProvider Protocol 的最小桩:记录调用,返回固定向量."""
|
||||
|
||||
model_name = "fake"
|
||||
version = "v1"
|
||||
dimensions = 4
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[list[str]] = []
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
self.calls.append(list(texts))
|
||||
return [[0.25, 0.25, 0.25, 0.25] for _ in texts]
|
||||
|
||||
|
||||
class FakeSession:
|
||||
def __init__(self) -> None:
|
||||
self.commits = 0
|
||||
self.rollbacks = 0
|
||||
|
||||
async def commit(self) -> None:
|
||||
self.commits += 1
|
||||
|
||||
async def rollback(self) -> None:
|
||||
self.rollbacks += 1
|
||||
|
||||
|
||||
class FakeClock:
|
||||
"""可手动拨动的单调时钟:TTL 断言不靠真 sleep,边界精确到微秒."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.now = 1000.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
def advance(self, seconds: float) -> None:
|
||||
self.now += seconds
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stubbed_repo(monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]:
|
||||
"""repo 的两个 IO 函数换成计数桩:``search_calls`` 就是「DB 交互次数」."""
|
||||
|
||||
calls: dict[str, Any] = {"search_calls": 0, "gaps": [], "results": []}
|
||||
|
||||
async def fake_hybrid_search(
|
||||
session: Any, query_vec: Any, query_text: str, **kwargs: Any
|
||||
) -> list[ScoredChunk]:
|
||||
calls["search_calls"] += 1
|
||||
return list(calls["results"])
|
||||
|
||||
async def fake_record_gap(session: Any, query_text: str, filter_ctx: Any) -> None:
|
||||
calls["gaps"].append((query_text, filter_ctx))
|
||||
|
||||
monkeypatch.setattr(repo, "hybrid_search", fake_hybrid_search)
|
||||
monkeypatch.setattr(repo, "record_gap", fake_record_gap)
|
||||
return calls
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _hermetic_module_state() -> Iterator[None]:
|
||||
"""用例结束还原模块级全局(失效代数、共享缓存实例),不污染同进程其它用例."""
|
||||
saved_version = repo.INVALIDATION_VERSION
|
||||
yield
|
||||
repo.INVALIDATION_VERSION = saved_version
|
||||
service_mod._shared_cache = None
|
||||
|
||||
|
||||
def _scored(chunk_id: int = 1) -> ScoredChunk:
|
||||
return ScoredChunk(
|
||||
chunk_id=chunk_id,
|
||||
entity_id=100,
|
||||
chunk_type="text",
|
||||
text="chunk",
|
||||
vec_score=1.0,
|
||||
kw_score=0.0,
|
||||
rule_score=0.0,
|
||||
freshness_score=1.0,
|
||||
final_score=0.55,
|
||||
)
|
||||
|
||||
|
||||
def _key(query_text: str = "q", **overrides: Any) -> CacheKey:
|
||||
defaults: dict[str, Any] = {
|
||||
"version": 1,
|
||||
"query_text": query_text,
|
||||
"top_k": 5,
|
||||
"body_shape": None,
|
||||
"occasion": None,
|
||||
"entity_types": (),
|
||||
"rule_boost": None,
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return CacheKey.build(**defaults)
|
||||
|
||||
|
||||
def _svc(cache: TtlCache[Any] | None = None, **kwargs: Any) -> RetrievalService:
|
||||
return RetrievalService(FakeSession(), embedding=FakeEmbedder(), cache=cache, **kwargs)
|
||||
|
||||
|
||||
# --- 1. 同参数二连查:命中零 DB ------------------------------------------------
|
||||
async def test_repeated_search_hits_cache_with_zero_db(stubbed_repo: dict[str, Any]) -> None:
|
||||
stubbed_repo["results"] = [_scored(1), _scored(2)]
|
||||
embedder = FakeEmbedder()
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=60.0, max_entries=16)
|
||||
svc = RetrievalService(FakeSession(), embedding=embedder, cache=cache)
|
||||
|
||||
first = await svc.search("面试穿什么")
|
||||
second = await svc.search("面试穿什么")
|
||||
|
||||
assert first == second
|
||||
assert [c.chunk_id for c in second.results] == [1, 2]
|
||||
assert stubbed_repo["search_calls"] == 1 # 第二次零 DB
|
||||
assert embedder.calls == [["面试穿什么"]] # 第二次连 embed 都跳过
|
||||
assert cache.stats() == {"hits": 1, "misses": 1, "evictions": 0, "entries": 1}
|
||||
|
||||
# 返回值是缓存内容的拷贝:调用方原地改 list(哪怕改的是命中返回的对象)
|
||||
# 也不会污染缓存(红线 2)
|
||||
second.results.append(_scored(99))
|
||||
third = await svc.search("面试穿什么")
|
||||
assert [c.chunk_id for c in third.results] == [1, 2]
|
||||
|
||||
|
||||
# --- 2. 键分量矩阵:任一影响结果的输入变化 → miss -----------------------------
|
||||
@pytest.mark.parametrize(
|
||||
("label", "query", "ctx", "top_k"),
|
||||
[
|
||||
("query_text", "另一个完全不同的问题", None, None),
|
||||
("occasion", None, UserContext(occasion="wedding"), None),
|
||||
("body_shape", None, UserContext(body_shape="pear"), None),
|
||||
("entity_types", None, UserContext(entity_types=("item",)), None),
|
||||
("rule_boost", None, UserContext(rule_boost={7: 0.5}), None),
|
||||
("top_k", None, None, 10),
|
||||
],
|
||||
)
|
||||
async def test_any_key_component_change_is_a_miss(
|
||||
stubbed_repo: dict[str, Any],
|
||||
label: str,
|
||||
query: str | None,
|
||||
ctx: UserContext | None,
|
||||
top_k: int | None,
|
||||
) -> None:
|
||||
stubbed_repo["results"] = [_scored(1)]
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=60.0, max_entries=16)
|
||||
svc = _svc(cache)
|
||||
await svc.search("面试穿什么") # 填充缓存
|
||||
|
||||
if top_k is not None:
|
||||
# top_k 是 service 级参数:同缓存、不同 top_k 的另一个实例
|
||||
await _svc(cache, top_k=top_k).search("面试穿什么")
|
||||
else:
|
||||
await svc.search(query or "面试穿什么", user_ctx=ctx)
|
||||
|
||||
assert stubbed_repo["search_calls"] == 2, label # 键不同 → 重新检索
|
||||
assert cache.misses == 2 and cache.hits == 0, label
|
||||
assert cache.stats()["entries"] == 2, label # 两个键并存,互不覆盖
|
||||
|
||||
|
||||
# --- 3. TTL 过期 → miss(注入 clock)------------------------------------------
|
||||
async def test_ttl_expiry_is_a_miss_then_refill(stubbed_repo: dict[str, Any]) -> None:
|
||||
stubbed_repo["results"] = [_scored(1)]
|
||||
clock = FakeClock()
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=60.0, max_entries=16, clock=clock)
|
||||
svc = _svc(cache)
|
||||
|
||||
await svc.search("q")
|
||||
clock.advance(60.0) # 恰好到达 TTL 边界:>= ttl 视作过期
|
||||
await svc.search("q")
|
||||
assert stubbed_repo["search_calls"] == 2
|
||||
assert cache.misses == 2 and cache.hits == 0
|
||||
|
||||
clock.advance(59.9) # 重新填充后未到期 → 命中
|
||||
await svc.search("q")
|
||||
assert stubbed_repo["search_calls"] == 2
|
||||
assert cache.stats() == {"hits": 1, "misses": 2, "evictions": 0, "entries": 1}
|
||||
|
||||
|
||||
# --- 4. LRU 容量超限 → 最旧被逐、计数正确 --------------------------------------
|
||||
def test_lru_evicts_oldest_beyond_capacity() -> None:
|
||||
cache: TtlCache[int] = TtlCache(ttl_seconds=60.0, max_entries=2)
|
||||
cache.put(_key("a"), 1)
|
||||
cache.put(_key("b"), 2)
|
||||
cache.put(_key("c"), 3) # 超限:最旧的 a 被逐
|
||||
|
||||
assert cache.get(_key("a")) is None
|
||||
assert cache.get(_key("b")) == 2
|
||||
assert cache.stats() == {"hits": 1, "misses": 1, "evictions": 1, "entries": 2}
|
||||
|
||||
|
||||
def test_lru_recency_updates_on_hit() -> None:
|
||||
"""命中会把条目移到"最新"端:之后超限逐出的是另一个最旧条目."""
|
||||
cache: TtlCache[int] = TtlCache(ttl_seconds=60.0, max_entries=2)
|
||||
cache.put(_key("a"), 1)
|
||||
cache.put(_key("b"), 2)
|
||||
assert cache.get(_key("a")) == 1 # a 变为最新
|
||||
cache.put(_key("c"), 3) # 逐出的应是 b,不是 a
|
||||
|
||||
assert cache.get(_key("b")) is None
|
||||
assert cache.get(_key("a")) == 1
|
||||
assert cache.evictions == 1
|
||||
|
||||
|
||||
def test_ttl_cache_rejects_invalid_params() -> None:
|
||||
with pytest.raises(ValueError, match="ttl_seconds"):
|
||||
TtlCache(ttl_seconds=0, max_entries=8)
|
||||
with pytest.raises(ValueError, match="max_entries"):
|
||||
TtlCache(ttl_seconds=60.0, max_entries=0)
|
||||
|
||||
|
||||
# --- 5. 空结果不缓存:gap 埋点每次都跑 -----------------------------------------
|
||||
async def test_empty_results_are_not_cached(stubbed_repo: dict[str, Any]) -> None:
|
||||
stubbed_repo["results"] = []
|
||||
session = FakeSession()
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=60.0, max_entries=16)
|
||||
svc = RetrievalService(session, embedding=FakeEmbedder(), cache=cache)
|
||||
|
||||
first = await svc.search("不存在的知识点")
|
||||
second = await svc.search("不存在的知识点")
|
||||
|
||||
assert first.degraded and second.degraded
|
||||
assert stubbed_repo["search_calls"] == 2 # 每次都真检索
|
||||
assert len(stubbed_repo["gaps"]) == 2 # gap 埋点不被缓存吞掉
|
||||
assert cache.stats() == {"hits": 0, "misses": 2, "evictions": 0, "entries": 0}
|
||||
assert session.commits == 2 # 每次埋点都照常提交
|
||||
|
||||
|
||||
# --- 6. invalidate_cache() 后同参数 → miss(代数失效)-------------------------
|
||||
async def test_invalidate_cache_invalidates_by_generation(
|
||||
stubbed_repo: dict[str, Any],
|
||||
) -> None:
|
||||
stubbed_repo["results"] = [_scored(1)]
|
||||
cache: TtlCache[Any] = TtlCache(ttl_seconds=60.0, max_entries=16)
|
||||
svc = _svc(cache)
|
||||
|
||||
await svc.search("面试穿什么")
|
||||
await svc.search("面试穿什么")
|
||||
assert stubbed_repo["search_calls"] == 1 # 命中
|
||||
|
||||
service_mod.invalidate_cache()
|
||||
await svc.search("面试穿什么")
|
||||
assert stubbed_repo["search_calls"] == 2 # 代数失效 → 重新检索
|
||||
|
||||
await svc.search("面试穿什么")
|
||||
assert stubbed_repo["search_calls"] == 2 # 新代重新缓存,再次命中
|
||||
|
||||
|
||||
def test_invalidate_cache_bumps_generation(stubbed_repo: dict[str, Any]) -> None:
|
||||
"""自增 repo.INVALIDATION_VERSION 的既有语义保留(T12 依赖该契约)."""
|
||||
before = repo.INVALIDATION_VERSION
|
||||
service_mod.invalidate_cache()
|
||||
assert before + 1 == repo.INVALIDATION_VERSION
|
||||
|
||||
|
||||
# --- 7. cache=None → 行为与现在完全一致 ----------------------------------------
|
||||
async def test_cache_none_searches_db_every_time(stubbed_repo: dict[str, Any]) -> None:
|
||||
stubbed_repo["results"] = [_scored(1)]
|
||||
embedder = FakeEmbedder()
|
||||
svc = RetrievalService(FakeSession(), embedding=embedder) # 不传 cache
|
||||
|
||||
first = await svc.search("面试穿什么")
|
||||
second = await svc.search("面试穿什么")
|
||||
|
||||
assert first == second
|
||||
assert stubbed_repo["search_calls"] == 2
|
||||
assert len(embedder.calls) == 2
|
||||
|
||||
|
||||
async def test_cache_none_degraded_path_records_gap_every_time(
|
||||
stubbed_repo: dict[str, Any],
|
||||
) -> None:
|
||||
stubbed_repo["results"] = []
|
||||
svc = RetrievalService(FakeSession(), embedding=FakeEmbedder())
|
||||
|
||||
await svc.search("问不到")
|
||||
await svc.search("问不到")
|
||||
|
||||
assert stubbed_repo["search_calls"] == 2
|
||||
assert len(stubbed_repo["gaps"]) == 2
|
||||
|
||||
|
||||
# --- rule_boost 哈希:确定性(排序归一 + 不同值不折叠)------------------------
|
||||
def test_rule_boost_hash_is_insertion_order_insensitive() -> None:
|
||||
d1 = {1: 0.5, 2: 0.3}
|
||||
d2 = {2: 0.3}
|
||||
d2[1] = 0.5
|
||||
assert rule_boost_hash(d1) == rule_boost_hash(d2)
|
||||
|
||||
a = {1: 0.5}
|
||||
a[2] = 0.3
|
||||
b = {2: 0.3}
|
||||
b[1] = 0.5
|
||||
assert CacheKey.build(version=1, query_text="q", top_k=5, rule_boost=a) == (
|
||||
CacheKey.build(version=1, query_text="q", top_k=5, rule_boost=b)
|
||||
)
|
||||
|
||||
|
||||
def test_rule_boost_hash_distinguishes_values_and_empty() -> None:
|
||||
assert rule_boost_hash({}) != rule_boost_hash({1: 0.0})
|
||||
assert rule_boost_hash({1: 0.5}) != rule_boost_hash({1: 0.6})
|
||||
|
||||
|
||||
def test_entity_types_order_does_not_change_key() -> None:
|
||||
"""``= ANY(...)`` 与顺序无关:同集合不同顺序应共享同一条目."""
|
||||
assert _key(entity_types=("item", "style")) == _key(entity_types=("style", "item"))
|
||||
|
||||
|
||||
# --- from_settings 装配:按 Settings 开/关与参数化 -----------------------------
|
||||
def test_from_settings_disabled_yields_no_cache(stubbed_repo: dict[str, Any]) -> None:
|
||||
settings = Settings(_env_file=None, retrieval_cache_enabled=False)
|
||||
svc = RetrievalService.from_settings(FakeSession(), settings)
|
||||
assert svc._cache is None
|
||||
|
||||
|
||||
def test_from_settings_attaches_shared_cache_with_settings_params(
|
||||
stubbed_repo: dict[str, Any],
|
||||
) -> None:
|
||||
settings = Settings(
|
||||
_env_file=None,
|
||||
retrieval_cache_ttl_seconds=7,
|
||||
retrieval_cache_max_entries=9,
|
||||
)
|
||||
first = RetrievalService.from_settings(FakeSession(), settings)
|
||||
assert first._cache is not None
|
||||
assert first._cache.ttl_seconds == 7
|
||||
assert first._cache.max_entries == 9
|
||||
|
||||
# 每请求实例共享同一缓存:跨请求复用是缓存的意义所在
|
||||
second = RetrievalService.from_settings(FakeSession(), settings)
|
||||
assert second._cache is first._cache
|
||||
|
||||
# TTL/容量变化 → 按新配置重建(进程内配置恒定,重建只在测试/热改配置发生)
|
||||
changed = Settings(_env_file=None, retrieval_cache_ttl_seconds=1, retrieval_cache_max_entries=9)
|
||||
third = RetrievalService.from_settings(FakeSession(), changed)
|
||||
assert third._cache is not None
|
||||
assert third._cache is not first._cache
|
||||
assert third._cache.ttl_seconds == 1
|
||||
Reference in New Issue
Block a user