缓存策略¶
演示 Agent 系统中两个关键的基础设施模式: -
LRUCache:基于容量的 LRU 缓存(OrderedDict实现)。 -LLMResponseCache:按 prompt 哈希缓存 LLM 响应,支持 TTL。 -TokenBucketRateLimiter:令牌桶限流,保护下游 API。 全部基于标准库,便于离线测试。 对应教程:第11卷《系统设计》第1章 并发与缓存。
文件路径: examples/volume-11-system-design\caching.py
源代码¶
"""缓存与限流示例。
演示 Agent 系统中两个关键的基础设施模式:
- ``LRUCache``:基于容量的 LRU 缓存(``OrderedDict`` 实现)。
- ``LLMResponseCache``:按 prompt 哈希缓存 LLM 响应,支持 TTL。
- ``TokenBucketRateLimiter``:令牌桶限流,保护下游 API。
全部基于标准库,便于离线测试。
对应教程:第11卷《系统设计》第1章 并发与缓存。
"""
from __future__ import annotations
import hashlib
import time
from collections import OrderedDict
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
class LRUCache:
"""基于容量的 LRU 缓存。"""
def __init__(self, capacity: int = 128) -> None:
if capacity <= 0:
raise ValueError("capacity 必须为正")
self.capacity = capacity
self._data: OrderedDict[Any, Any] = OrderedDict()
self.hits = 0
self.misses = 0
def get(self, key: Any, default: Any = None) -> Any:
"""获取值,命中时移到末尾(最近使用)。"""
if key in self._data:
self._data.move_to_end(key)
self.hits += 1
return self._data[key]
self.misses += 1
return default
def set(self, key: Any, value: Any) -> None:
"""写入值,超容量时淘汰最久未使用的键。"""
if key in self._data:
self._data.move_to_end(key)
self._data[key] = value
if len(self._data) > self.capacity:
self._data.popitem(last=False)
def __contains__(self, key: Any) -> bool:
return key in self._data
def __len__(self) -> int:
return len(self._data)
def hash_prompt(prompt: str, model: str = "") -> str:
"""对 prompt + model 取哈希作为缓存键。
Args:
prompt: 提示文本。
model: 模型名。
Returns:
16 位哈希字符串。
"""
raw = f"{model}::{prompt}"
return hashlib.md5(raw.encode("utf-8")).hexdigest()[:16] # noqa: S324
@dataclass
class LLMResponseCache:
"""LLM 响应缓存,支持 TTL。
Attributes:
max_size: 最大条目数。
default_ttl: 默认存活秒数,``None`` 永不过期。
"""
max_size: int = 256
default_ttl: float | None = 300
_store: dict[str, tuple[Any, float | None]] = field(default_factory=dict)
def _now(self) -> float:
return time.monotonic()
def get(self, prompt: str, model: str = "default") -> Any:
"""获取缓存的响应。"""
key = hash_prompt(prompt, model)
if key not in self._store:
return None
value, expire_at = self._store[key]
if expire_at is not None and self._now() >= expire_at:
del self._store[key]
return None
return value
def set(self, prompt: str, response: Any, model: str = "default", ttl: float | None = "default") -> None: # type: ignore[assignment]
"""写入响应。"""
if ttl == "default":
ttl = self.default_ttl
key = hash_prompt(prompt, model)
expire_at = (self._now() + ttl) if ttl is not None else None
self._store[key] = (response, expire_at)
if len(self._store) > self.max_size:
# 淘汰最旧(按插入顺序)
oldest = next(iter(self._store))
del self._store[oldest]
def cached_call(
self,
prompt: str,
llm: Callable[[str], Any],
model: str = "default",
) -> Any:
"""带缓存的 LLM 调用:命中则直接返回,否则调用 ``llm`` 并缓存。"""
cached = self.get(prompt, model)
if cached is not None:
return cached
response = llm(prompt)
self.set(prompt, response, model)
return response
@dataclass
class TokenBucketRateLimiter:
"""令牌桶限流器。
Attributes:
capacity: 桶容量(最大突发请求数)。
refill_rate: 每秒补充的令牌数。
"""
capacity: float
refill_rate: float
_tokens: float = field(init=False)
_last_refill: float = field(init=False)
def __post_init__(self) -> None:
if self.capacity <= 0 or self.refill_rate <= 0:
raise ValueError("capacity 与 refill_rate 必须为正")
self._tokens = self.capacity
self._last_refill = time.monotonic()
def _refill(self) -> None:
now = time.monotonic()
elapsed = now - self._last_refill
self._tokens = min(self.capacity, self._tokens + elapsed * self.refill_rate)
self._last_refill = now
def acquire(self, tokens: float = 1.0) -> bool:
"""尝试获取 ``tokens`` 个令牌。
Args:
tokens: 需要的令牌数。
Returns:
成功获取返回 True,否则 False(不阻塞)。
"""
if tokens <= 0:
raise ValueError("tokens 必须为正")
self._refill()
if self._tokens >= tokens:
self._tokens -= tokens
return True
return False
@property
def available_tokens(self) -> float:
"""当前可用令牌数(补充后)。"""
self._refill()
return self._tokens
def main() -> None:
"""演示缓存与限流。"""
print("== LRU 缓存 ==")
cache = LRUCache(capacity=2)
cache.set("a", 1)
cache.set("b", 2)
cache.get("a") # a 最近使用
cache.set("c", 3) # 淘汰 b
print(" contains b:", "b" in cache, " contains a:", "a" in cache)
print(f" hits={cache.hits} misses={cache.misses}")
print("\n== LLM 响应缓存 ==")
llm_cache = LLMResponseCache(default_ttl=100)
call_count = 0
def fake_llm(prompt: str) -> str:
nonlocal call_count
call_count += 1
return f"reply:{prompt}"
print(" 第一次:", llm_cache.cached_call("你好", fake_llm))
print(" 第二次:", llm_cache.cached_call("你好", fake_llm))
print(f" 实际调用次数: {call_count}")
print("\n== 限流器 ==")
rl = TokenBucketRateLimiter(capacity=3, refill_rate=1)
results = [rl.acquire() for _ in range(5)]
print(f" 连续 5 次获取: {results}")
if __name__ == "__main__":
main()