import hashlib import logging import pickle import threading import time from collections import OrderedDict from typing import Any, Protocol, cast from injector import inject, singleton from private_gpt.settings.settings import Settings from private_gpt.utils.dependencies import format_missing_dependency_message logger = logging.getLogger(__name__) class Cache(Protocol): def get(self, namespace: str, key: str) -> Any | None: ... def set( self, namespace: str, key: str, value: Any, ttl_seconds: int | None = None, ) -> None: ... def delete(self, namespace: str, key: str) -> None: ... class MemoryCache: def __init__(self, max_entries: int = 1000) -> None: self._max_entries = max_entries self._values: OrderedDict[str, tuple[float | None, Any]] = OrderedDict() self._lock = threading.RLock() @staticmethod def _key(namespace: str, key: str) -> str: digest = hashlib.sha256(key.encode()).hexdigest() return f"{namespace}:{digest}" def get(self, namespace: str, key: str) -> Any | None: cache_key = self._key(namespace, key) with self._lock: entry = self._values.get(cache_key) if entry is None: return None expires_at, value = entry if expires_at is not None and expires_at <= time.monotonic(): del self._values[cache_key] return None self._values.move_to_end(cache_key) return value def set( self, namespace: str, key: str, value: Any, ttl_seconds: int | None = None, ) -> None: cache_key = self._key(namespace, key) expires_at = time.monotonic() + ttl_seconds if ttl_seconds is not None else None with self._lock: self._values[cache_key] = (expires_at, value) self._values.move_to_end(cache_key) while len(self._values) > self._max_entries: self._values.popitem(last=False) def delete(self, namespace: str, key: str) -> None: with self._lock: self._values.pop(self._key(namespace, key), None) class RedisCache: def __init__(self, settings: Settings) -> None: try: import redis except ImportError as error: raise ImportError( format_missing_dependency_message("Redis cache", extras="redis") ) from error redis_database = ( settings.cache.redis_database if settings.cache.redis_database is not None else int(settings.redis.database or 0) ) self._prefix = settings.cache.key_prefix self._client = redis.Redis.from_url( settings.redis.url, db=redis_database, decode_responses=False, ) def _key(self, namespace: str, key: str) -> str: digest = hashlib.sha256(key.encode()).hexdigest() return f"{self._prefix}:{namespace}:{digest}" def get(self, namespace: str, key: str) -> Any | None: value = self._client.get(self._key(namespace, key)) return pickle.loads(cast(bytes, value)) if value is not None else None def set( self, namespace: str, key: str, value: Any, ttl_seconds: int | None = None, ) -> None: self._client.set( self._key(namespace, key), pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL), ex=ttl_seconds, ) def delete(self, namespace: str, key: str) -> None: self._client.delete(self._key(namespace, key)) def close(self) -> None: self._client.close() # type: ignore[no-untyped-call] @singleton class CacheService: @inject def __init__(self, settings: Settings) -> None: self._default_ttl = settings.cache.ttl_seconds self._memory = MemoryCache(settings.cache.max_entries) self._distributed = ( RedisCache(settings) if settings.cache.provider == "redis" else None ) def get(self, namespace: str, key: str) -> Any | None: value = self._memory.get(namespace, key) if value is not None or self._distributed is None: return value try: value = self._distributed.get(namespace, key) except Exception: logger.warning("Distributed cache read failed", exc_info=True) return None if value is not None: self._memory.set(namespace, key, value, self._default_ttl) return value def set( self, namespace: str, key: str, value: Any, ttl_seconds: int | None = None, ) -> None: effective_ttl = self._default_ttl if ttl_seconds is None else ttl_seconds self._memory.set(namespace, key, value, effective_ttl) if self._distributed is not None: try: self._distributed.set(namespace, key, value, effective_ttl) except Exception: logger.warning("Distributed cache write failed", exc_info=True) def delete(self, namespace: str, key: str) -> None: self._memory.delete(namespace, key) if self._distributed is not None: try: self._distributed.delete(namespace, key) except Exception: logger.warning("Distributed cache delete failed", exc_info=True) def close(self) -> None: if self._distributed is not None: self._distributed.close()