File size: 1,711 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """In-memory snapshot cache."""
from __future__ import annotations
import hashlib
import logging
import time
from collections import OrderedDict
from aistudio_api.config import settings
logger = logging.getLogger("aistudio")
class SnapshotCache:
def __init__(self, ttl: int | None = None, max_size: int | None = None):
self._cache: OrderedDict[str, tuple] = OrderedDict()
self.ttl = ttl or settings.snapshot_cache_ttl
self.max_size = max_size or settings.snapshot_cache_max
@staticmethod
def _hash(prompt: str) -> str:
return hashlib.sha256(prompt.encode()).hexdigest()
def get(self, prompt: str) -> tuple | None:
key = self._hash(prompt)
if key not in self._cache:
return None
snapshot, url, headers, body, ts = self._cache[key]
age = time.time() - ts
if age >= self.ttl:
del self._cache[key]
logger.info("Snapshot 缓存过期: %s...", key[:8])
return None
# LRU: move to end
self._cache.move_to_end(key)
logger.info("Snapshot 缓存命中: %s... (%ss ago)", key[:8], int(age))
return snapshot, url, headers, body
def put(self, prompt: str, snapshot: str, url: str, headers: dict, body: str):
key = self._hash(prompt)
# Evict oldest if at capacity
while len(self._cache) >= self.max_size:
evicted_key, _ = self._cache.popitem(last=False)
logger.info("Snapshot 缓存淘汰: %s...", evicted_key[:8])
self._cache[key] = (snapshot, url, headers, body, time.time())
logger.info("Snapshot 已缓存: %s...", key[:8])
def clear(self):
self._cache.clear()
|