Spaces:
Running
Running
File size: 4,377 Bytes
167c070 da3be01 167c070 da3be01 167c070 da3be01 167c070 da3be01 167c070 da3be01 167c070 da3be01 167c070 da3be01 167c070 da3be01 167c070 | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """Cache backends. Redis when available, otherwise in-process cachetools.TTLCache."""
import asyncio
import hashlib
import json
import logging
import os
from typing import Any, Optional
logger = logging.getLogger("rga_auditor.cache")
def _redis_url() -> str:
return os.getenv("REDIS_URL", "")
CACHE_TTL = {
"llm_response": int(os.getenv("CACHE_TTL_LLM", "3600")),
"embedding": int(os.getenv("CACHE_TTL_EMBEDDING", "86400")),
"document": int(os.getenv("CACHE_TTL_DOCUMENT", "300")),
"query_result": int(os.getenv("CACHE_TTL_QUERY", "600")),
"user": int(os.getenv("CACHE_TTL_USER", "60")),
}
DEFAULT_MAX_ENTRIES = int(os.getenv("CACHE_MAX_ENTRIES", "10000"))
class CacheBackend:
async def get(self, key: str) -> Optional[Any]: ...
async def set(self, key: str, value: Any, ttl: int) -> None: ...
async def delete(self, key: str) -> None: ...
async def flush_pattern(self, pattern: str) -> None: ...
class MemoryBackend(CacheBackend):
def __init__(self, max_size: int = DEFAULT_MAX_ENTRIES) -> None:
self._store: dict[str, Any] = {}
self._expires: dict[str, float] = {}
self._max_size = max_size
async def get(self, key: str) -> Optional[Any]:
self._evict_expired()
try:
return self._store[key]
except KeyError:
return None
async def set(self, key: str, value: Any, ttl: int) -> None:
self._evict_expired()
if len(self._store) >= self._max_size:
return
import time
self._store[key] = value
self._expires[key] = time.monotonic() + ttl
async def delete(self, key: str) -> None:
self._store.pop(key, None)
self._expires.pop(key, None)
async def flush_pattern(self, pattern: str) -> None:
for k in [k for k in list(self._store.keys()) if pattern in k]:
self._store.pop(k, None)
self._expires.pop(k, None)
def _evict_expired(self) -> None:
import time
now = time.monotonic()
expired = [k for k, t in self._expires.items() if t < now]
for k in expired:
self._store.pop(k, None)
self._expires.pop(k, None)
class RedisBackend(CacheBackend):
def __init__(self, url: str) -> None:
import redis.asyncio as redis
self._client = redis.from_url(url, decode_responses=True)
async def get(self, key: str) -> Optional[Any]:
val = await self._client.get(key)
if val is None:
return None
try:
return json.loads(val)
except (json.JSONDecodeError, TypeError):
return val
async def set(self, key: str, value: Any, ttl: int) -> None:
payload = json.dumps(value, default=str)
await self._client.setex(key, ttl, payload)
async def delete(self, key: str) -> None:
await self._client.delete(key)
async def flush_pattern(self, pattern: str) -> None:
cursor = 0
while True:
cursor, keys = await self._client.scan(cursor, match=f"*{pattern}*", count=100)
if keys:
await self._client.delete(*keys)
if cursor == 0:
break
def _build_backend() -> CacheBackend:
url = _redis_url()
if url:
try:
backend = RedisBackend(url)
logger.info("Cache backend: Redis (%s)", url.split("@")[-1] if "@" in url else "local")
return backend
except Exception as e:
logger.warning("Redis unavailable (%s) — falling back to in-memory cache", e)
logger.info("Cache backend: in-memory (TTLCache, max=%d)", DEFAULT_MAX_ENTRIES)
return MemoryBackend()
_cache: Optional[CacheBackend] = None
def get_cache() -> CacheBackend:
global _cache
if _cache is None:
_cache = _build_backend()
return _cache
def cache_key(prefix: str, *parts: str) -> str:
raw = ":".join(parts)
h = hashlib.sha256(raw.encode()).hexdigest()
return f"{prefix}:{h}"
async def get_or_compute(key: str, ttl: int, fn):
cache = get_cache()
cached = await cache.get(key)
if cached is not None:
return cached
if callable(fn):
result = fn()
value = await result if asyncio.iscoroutine(result) else result
else:
value = fn
await cache.set(key, value, ttl)
return value
|