Spaces:
Running
Running
File size: 2,659 Bytes
fb0ec21 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Optional Upstash Redis cache (REST) with an in-process TTL fallback.
Upstash credentials arrive as REST URL + token (the console shows both). When
they are absent the engine uses a process-local TTL dict — correct per
container, lost on restart, which is acceptable for enrichment caching.
"""
import json
import logging
import time
from typing import Any, Optional
import httpx
from ..config import settings
logger = logging.getLogger("abm.cache")
class Cache:
def __init__(self) -> None:
self._http: Optional[httpx.AsyncClient] = None
self._mem: dict[str, tuple[float, str]] = {}
@property
def mode(self) -> str:
return "upstash" if settings.upstash_configured else "memory"
def _client(self) -> httpx.AsyncClient:
if self._http is None:
self._http = httpx.AsyncClient(
base_url=settings.upstash_rest_url.rstrip("/"),
headers={"Authorization": f"Bearer {settings.upstash_rest_token}"},
timeout=httpx.Timeout(10.0),
)
return self._http
async def get(self, key: str) -> Optional[Any]:
if settings.upstash_configured:
try:
res = await self._client().get(f"/get/{key}")
if res.status_code == 200:
raw = res.json().get("result")
return json.loads(raw) if raw else None
except Exception as exc:
logger.warning("cache get failed: %s", exc)
return None
item = self._mem.get(key)
if not item:
return None
expires, raw = item
if expires < time.time():
self._mem.pop(key, None)
return None
return json.loads(raw)
async def set(self, key: str, value: Any, ttl_s: int = 3600) -> None:
raw = json.dumps(value, default=str)
if settings.upstash_configured:
try:
await self._client().get(f"/set/{key}?EX={int(ttl_s)}",
headers={"Content-Type": "text/plain"},
content=raw.encode())
return
except Exception as exc:
logger.warning("cache set failed: %s", exc)
return
self._mem[key] = (time.time() + ttl_s, raw)
async def ping(self) -> bool:
if not settings.upstash_configured:
return True # memory mode always "works"
try:
res = await self._client().get("/ping")
return res.status_code == 200
except Exception:
return False
cache = Cache()
|