feat: add response cache middleware with poisoning-safe keys

This commit is contained in:
2026-07-20 07:10:42 -04:00
parent c3d5079d39
commit 6286086551
4 changed files with 434 additions and 0 deletions
+30
View File
@@ -0,0 +1,30 @@
"""进程内缓存后端: dict + 过期时刻;测试与无 Redis 场景使用。"""
from __future__ import annotations
import time
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Callable
class InMemoryCache:
"""笨 KV(CacheBackend 契约);容量不设限——库内默认 TTL 必填防无界增长。"""
def __init__(self, now: Callable[[], float] = time.monotonic) -> None:
self._now = now
self._store: dict[str, tuple[float, str]] = {}
async def get(self, key: str) -> str | None:
item = self._store.get(key)
if item is None:
return None
expires_at, value = item
if self._now() >= expires_at:
del self._store[key]
return None
return value
async def set(self, key: str, value: str, ttl_s: int) -> None:
self._store[key] = (self._now() + ttl_s, value)
+44
View File
@@ -0,0 +1,44 @@
"""Redis 缓存后端: 笨 KV(CacheBackend 契约);redis 是 optional extra。
降级不在这里做——本类忠实上抛异常,静默降级是 CacheMW 的职责(算法一份)。
"""
from __future__ import annotations
from typing import Any
try:
import redis.asyncio as aioredis
except ImportError as _exc: # pragma: no cover - 依赖缺失路径
aioredis = None
_IMPORT_ERROR = _exc
else:
_IMPORT_ERROR = None
class RedisCache:
"""基于 redis.asyncio 的 KV;value 为 JSON 字符串,TTL 由调用方传入。"""
def __init__(self, client: Any) -> None:
if aioredis is None:
raise ImportError(
"Redis 缓存后端需要 redis 包: pip install 'polygateway[redis]'"
) from _IMPORT_ERROR
self._client = client
@classmethod
def from_url(cls, url: str) -> RedisCache:
if aioredis is None:
raise ImportError(
"Redis 缓存后端需要 redis 包: pip install 'polygateway[redis]'"
) from _IMPORT_ERROR
return cls(aioredis.from_url(url, decode_responses=True))
async def get(self, key: str) -> str | None:
return await self._client.get(key)
async def set(self, key: str, value: str, ttl_s: int) -> None:
await self._client.set(key, value, ex=ttl_s)
async def aclose(self) -> None:
await self._client.aclose()
+155
View File
@@ -0,0 +1,155 @@
"""CacheMW: 响应缓存中间件;key 公式在此(算法一份),后端是笨 KV。
ARCH §7.5: key = sha256(canonical_json({model, messages_digest, namespace,
salt}));多模态 part 先摘要再 hash(防 Video-Tree 整段 base64 进 hash 的
开销);namespace 必填防跨项目/租户毒化;只缓存阶梯通过的成功响应
(StructuredMW 在内层,能返回即已通过)。读写失败静默降级(铁律)。
"""
from __future__ import annotations
import asyncio
import dataclasses
import hashlib
import json
import uuid
from typing import TYPE_CHECKING, Any
from loguru import logger
from polygateway.types import ChatRequest, LLMResponse
if TYPE_CHECKING:
from polygateway.ports import CacheBackend, CallNext, StructuredOutputStrategy
_KEY_PREFIX = "pgw:cache:"
_RESPONSE_FIELDS = {f.name for f in dataclasses.fields(LLMResponse)}
def digest_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""多模态 content part 先各自 sha256 摘要再参与序列化;文本原文参与。
与遥测落库共用同一函数(ARCH §7.8),保证缓存 key 与遥测口径一致。
"""
digested = []
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
parts = [_digest_part(part) for part in content]
digested.append({**msg, "content": parts})
else:
digested.append(msg)
return digested
def _digest_part(part: Any) -> Any:
if isinstance(part, dict) and part.get("type") == "image_url":
url = str(part.get("image_url", {}).get("url", ""))
return {"type": "image_url", "sha256": hashlib.sha256(url.encode()).hexdigest()}
return part
def build_cache_key(
model_fingerprint: str, messages: list[dict[str, Any]], namespace: str, salt: str | None
) -> str:
"""缓存 key 公式;salt 仅非 None 时参与(VT 旧键语义: 不传 salt 键形不变)。"""
key_obj: dict[str, Any] = {
"model": model_fingerprint,
"messages": digest_messages(messages),
"namespace": namespace,
}
if salt is not None:
key_obj["salt"] = salt
payload = json.dumps(key_obj, sort_keys=True, ensure_ascii=False)
return _KEY_PREFIX + hashlib.sha256(payload.encode("utf-8")).hexdigest()
class CacheMW:
"""洋葱第二层(遥测内、结构化外)。
model_fingerprint 由装配层从源列表计算(多源 scope = 排序去重的 model
名合集);源集合变化 → key 变化 → 一次性冷启动,换取跨源集合零毒化。
"""
def __init__(
self,
*,
backend: CacheBackend,
model_fingerprint: str,
default_namespace: str,
ttl_s: int,
strategy: StructuredOutputStrategy | None = None,
) -> None:
if not default_namespace.strip():
raise ValueError("缓存 namespace 不能为空(防跨项目/租户毒化)")
if ttl_s <= 0:
raise ValueError("缓存 TTL 必须 > 0(禁止永不过期)")
self._backend = backend
self._fingerprint = model_fingerprint
self._namespace = default_namespace
self._ttl_s = ttl_s
self._strategy = strategy
async def __call__(self, request: ChatRequest, call_next: CallNext) -> LLMResponse:
namespace = request.cache_namespace or self._namespace
key = build_cache_key(self._fingerprint, request.messages, namespace, request.cache_salt)
cached = await self._safe_get(key)
if cached is not None:
hit = self._rehydrate(cached, request)
if hit is not None:
return hit
response = await call_next(request)
await self._safe_set(key, self._serialize(response))
return response
# —— 命中路径 ——
def _rehydrate(self, raw: str, request: ChatRequest) -> LLMResponse | None:
"""反序列化 + 按本次调用的 structured 档零网络重建;失败按未命中。"""
try:
data = json.loads(raw)
fields = {k: v for k, v in data.items() if k in _RESPONSE_FIELDS}
structured_data = self._rebuild_structured(fields.get("content", ""), request)
fields.update(
cache_hit=True, latency_ms=0, ttft_ms=None, max_inter_token_ms=None,
call_id=str(uuid.uuid4()), structured_data=structured_data,
)
return LLMResponse(**fields)
except Exception as exc:
logger.warning("缓存命中重建失败,按未命中回源: {}", exc)
return None
def _rebuild_structured(self, content: str, request: ChatRequest) -> Any | None:
"""对缓存 content 重跑阶梯②③(schema 变更后旧缓存自动重校验,设计 §2.1)。"""
if request.structured is None:
return None
if self._strategy is None:
raise ValueError("structured 调用需要装配 strategy 才能命中重建")
parsed = self._strategy.parse(content)
if request.structured == "json":
return parsed
return request.structured.model_validate(parsed)
# —— 写路径与降级 ——
def _serialize(self, response: LLMResponse) -> str:
data = dataclasses.asdict(response)
data.pop("structured_data", None) # pydantic 实例不可 JSON 往返(设计 §2.1)
return json.dumps(data, ensure_ascii=False)
async def _safe_get(self, key: str) -> str | None:
try:
return await self._backend.get(key)
except asyncio.CancelledError:
raise
except Exception as exc:
logger.warning("缓存读取失败,降级为未命中: {}", exc)
return None
async def _safe_set(self, key: str, value: str) -> None:
try:
await self._backend.set(key, value, self._ttl_s)
except asyncio.CancelledError:
raise
except Exception as exc:
logger.warning("缓存写入失败,跳过缓存: {}", exc)
+205
View File
@@ -0,0 +1,205 @@
"""CacheMW 与缓存 key 公式测试(ARCH §7.5: 防毒化 key、命中重建、静默降级)。"""
import dataclasses
import hashlib
import json
import pytest
from polygateway.backends.memory.cache import InMemoryCache
from polygateway.errors import ResultInvalidError, TransientError
from polygateway.middleware.cache import CacheMW, build_cache_key, digest_messages
from polygateway.types import ChatRequest, LLMResponse
_MSGS = [{"role": "user", "content": "hi"}]
def _resp(content="cached", **overrides):
base = dict(
content=content, thinking="", model="m", provider="p",
prompt_tokens=1, completion_tokens=2, latency_ms=30,
ttft_ms=5.0, max_inter_token_ms=2.0, cache_hit=False, call_id="orig",
source_name="s1", usage_source="measured",
)
base.update(overrides)
return LLMResponse(**base)
class TestKeyFormula:
def test_same_request_same_key(self):
k1 = build_cache_key("m", _MSGS, "proj", None)
k2 = build_cache_key("m", _MSGS, "proj", None)
assert k1 == k2 and k1.startswith("pgw:cache:")
@pytest.mark.parametrize(
("a", "b"),
[
(("m1", _MSGS, "proj", None), ("m2", _MSGS, "proj", None)),
(("m", _MSGS, "proj", None), ("m", _MSGS, "tenant2", None)),
(("m", _MSGS, "proj", None), ("m", _MSGS, "proj", "epoch2")),
(("m", _MSGS, "proj", "s1"), ("m", _MSGS, "proj", "s2")),
(("m", _MSGS, "proj", None), ("m", [{"role": "user", "content": "yo"}], "proj", None)),
],
)
def test_any_dimension_change_changes_key(self, a, b):
assert build_cache_key(*a) != build_cache_key(*b)
def test_multimodal_part_digested_not_inlined(self):
big_b64 = "data:image/png;base64," + "A" * 1_000_000
messages = [{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": big_b64}},
{"type": "text", "text": "describe"},
]}]
digested = digest_messages(messages)
payload = json.dumps(digested, ensure_ascii=False)
assert len(payload) < 500 # 大图不进 canonical_json
expected = hashlib.sha256(big_b64.encode()).hexdigest()
assert expected in payload # 但字节变化仍改变 key
# 图像字节变化 → key 变
messages2 = [{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": big_b64[:-1] + "B"}},
{"type": "text", "text": "describe"},
]}]
assert build_cache_key("m", messages, "p", None) != build_cache_key("m", messages2, "p", None)
class _Terminal:
def __init__(self, response):
self.response = response
self.calls = 0
async def __call__(self, request):
self.calls += 1
if isinstance(self.response, Exception):
raise self.response
return self.response
def _mw(backend, **kwargs):
defaults = dict(backend=backend, model_fingerprint="m", default_namespace="proj", ttl_s=3600)
defaults.update(kwargs)
return CacheMW(**defaults)
class TestCacheFlow:
async def test_miss_then_hit_with_fresh_call_id(self):
backend = InMemoryCache()
mw = _mw(backend)
terminal = _Terminal(_resp())
first = await mw(ChatRequest(messages=_MSGS), terminal)
assert first.cache_hit is False and terminal.calls == 1
second = await mw(ChatRequest(messages=_MSGS), terminal)
assert second.cache_hit is True and second.latency_ms == 0
assert second.content == "cached"
assert second.call_id != first.call_id # 命中生成独立 cache_call_id
assert terminal.calls == 1 # 未再触达内层
async def test_per_call_namespace_overrides_default(self):
backend = InMemoryCache()
mw = _mw(backend)
terminal = _Terminal(_resp())
await mw(ChatRequest(messages=_MSGS, cache_namespace="tenant-a"), terminal)
# 另一租户不得命中
await mw(ChatRequest(messages=_MSGS, cache_namespace="tenant-b"), terminal)
assert terminal.calls == 2
async def test_failure_not_cached(self):
backend = InMemoryCache()
mw = _mw(backend)
failing = _Terminal(TransientError("boom"))
with pytest.raises(TransientError):
await mw(ChatRequest(messages=_MSGS), failing)
ok = _Terminal(_resp())
await mw(ChatRequest(messages=_MSGS), ok)
assert ok.calls == 1 # 失败未被固化,正常回源
async def test_ttl_expiry(self):
t = {"now": 0.0}
backend = InMemoryCache(now=lambda: t["now"])
mw = _mw(backend, ttl_s=100)
terminal = _Terminal(_resp())
await mw(ChatRequest(messages=_MSGS), terminal)
t["now"] = 101.0
await mw(ChatRequest(messages=_MSGS), terminal)
assert terminal.calls == 2
class _BrokenBackend:
async def get(self, key):
raise ConnectionError("redis down")
async def set(self, key, value, ttl_s):
raise ConnectionError("redis down")
class TestDegradation:
async def test_backend_failure_degrades_silently(self):
mw = _mw(_BrokenBackend())
terminal = _Terminal(_resp())
resp = await mw(ChatRequest(messages=_MSGS), terminal)
assert resp.content == "cached" and terminal.calls == 1 # 读写全降级,调用照常
async def test_corrupt_cache_value_treated_as_miss(self):
backend = InMemoryCache()
mw = _mw(backend)
key = build_cache_key("m", _MSGS, "proj", None)
await backend.set(key, "{not json", 3600)
terminal = _Terminal(_resp())
resp = await mw(ChatRequest(messages=_MSGS), terminal)
assert terminal.calls == 1 and resp.cache_hit is False
class _FakeStrategy:
"""fake StructuredOutputStrategy(T2 冻结的 Protocol,不依赖 T11)。"""
def request_overlay(self, schema):
return {}
def parse(self, text):
data = json.loads(text) # 简化: 直接 json.loads
if not isinstance(data, dict):
raise ResultInvalidError("非对象", raw_text=text)
return data
class _StrictModel:
"""鸭子型 pydantic 模型: model_validate 要求含 answer 键。"""
@classmethod
def model_validate(cls, data):
if "answer" not in data:
raise ValueError("missing answer")
return {"validated": data["answer"]}
class TestStructuredRehydration:
async def test_hit_rebuilds_structured_data(self):
backend = InMemoryCache()
mw = _mw(backend, strategy=_FakeStrategy())
terminal = _Terminal(_resp(content='{"answer": 42}'))
req = ChatRequest(messages=_MSGS, structured=_StrictModel)
await mw(req, terminal)
hit = await mw(req, terminal)
assert hit.cache_hit is True
assert hit.structured_data == {"validated": 42}
assert terminal.calls == 1
async def test_schema_change_revalidation_failure_falls_back_to_source(self):
backend = InMemoryCache()
mw = _mw(backend, strategy=_FakeStrategy())
terminal = _Terminal(_resp(content='{"other": 1}')) # 缓存内容不含 answer
await mw(ChatRequest(messages=_MSGS), terminal) # 无 structured 写入
# 换 schema 读: 重校验失败 → 按未命中回源
again = await mw(ChatRequest(messages=_MSGS, structured=_StrictModel), terminal)
assert terminal.calls == 2 and again.cache_hit is False
async def test_structured_data_not_serialized_into_cache(self):
backend = InMemoryCache()
mw = _mw(backend, strategy=_FakeStrategy())
terminal = _Terminal(
dataclasses.replace(_resp(content='{"answer": 1}'), structured_data={"x": object()})
)
await mw(ChatRequest(messages=_MSGS), terminal) # 不可 JSON 的 structured_data 不阻塞写缓存
key = build_cache_key("m", _MSGS, "proj", None)
raw = await backend.get(key)
assert raw is not None and "structured_data" not in json.loads(raw)