跳转至

缓存策略

演示 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()