File size: 2,070 Bytes
9fb692f | 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 | """A minimal in-memory TTL cache, used to avoid re-fetching the
(effectively static) molecule catalog on every tool call that needs to
resolve a molecule name, while still eventually picking up a real catalog
change (e.g. after a kernel restart with a different build) instead of
caching forever like the pre-Phase-2 code did (a plain dict, reset only by
the test suite, never by production code). Also caches FAILURES for a
short TTL, so a tight loop of tool calls made while the kernel is down
doesn't each wait out their own full connection/timeout error -- see
`set_failure`/`get`'s re-raise behavior below.
"""
import time
from typing import Any, Optional
class TTLCache:
def __init__(self, ttl_seconds: float = 300.0, failure_ttl_seconds: float = 10.0):
self._ttl = ttl_seconds
self._failure_ttl = failure_ttl_seconds
# key -> (expires_at, value_or_exception, is_failure)
self._store: dict[str, tuple[float, Any, bool]] = {}
def get(self, key: str):
"""Returns the cached value, or None if missing/expired. Re-raises
a cached failure (so callers see the same error shape they would
from a fresh call) instead of returning it as if it were a value."""
entry = self._store.get(key)
if entry is None:
return None
expires_at, value, is_failure = entry
if time.monotonic() >= expires_at:
del self._store[key]
return None
if is_failure:
raise value
return value
def set(self, key: str, value: Any) -> None:
self._store[key] = (time.monotonic() + self._ttl, value, False)
def set_failure(self, key: str, exc: Exception) -> None:
self._store[key] = (time.monotonic() + self._failure_ttl, exc, True)
def invalidate(self, key: Optional[str] = None) -> None:
"""Clears one key, or the whole cache if key is None."""
if key is None:
self._store.clear()
else:
self._store.pop(key, None)
|