cache.py

"""
Caching Layer for Production RAG
Redis-based caching with TTL and automatic invalidation
"""

import json
import hashlib
import pickle
from typing import Any, Optional, List
from datetime import timedelta
import logging
from functools import wraps
import asyncio

import redis
from redis import asyncio as aioredis

from config import CACHE_CONFIG, REDIS_URL

logger = logging.getLogger(__name__)


class CacheKeyBuilder:
    """Build consistent cache keys"""
    
    @staticmethod
    def query_key(query: str, params: dict = None) -> str:
        """Generate cache key for query results"""
        key_data = {"query": query, "params": params or {}}
        key_hash = hashlib.md5(json.dumps(key_data, sort_keys=True).encode()).hexdigest()
        return f"rag:query:{key_hash}"
    
    @staticmethod
    def embedding_key(text: str) -> str:
        """Generate cache key for embeddings"""
        text_hash = hashlib.md5(text.encode()).hexdigest()
        return f"rag:embedding:{text_hash}"
    
    @staticmethod
    def document_key(doc_id: str) -> str:
        """Generate cache key for documents"""
        return f"rag:doc:{doc_id}"
    
    @staticmethod
    def search_key(query: str, top_k: int, filters: dict = None) -> str:
        """Generate cache key for search results"""
        key_data = {"query": query, "top_k": top_k, "filters": filters or {}}
        key_hash = hashlib.md5(json.dumps(key_data, sort_keys=True).encode()).hexdigest()
        return f"rag:search:{key_hash}"


class RedisCache:
    """Synchronous Redis cache implementation"""
    
    def __init__(self, url: str = REDIS_URL):
        self.enabled = CACHE_CONFIG["enabled"]
        self.default_ttl = CACHE_CONFIG["ttl_seconds"]
        self.embedding_ttl = CACHE_CONFIG["embedding_cache_ttl"]
        
        if self.enabled:
            try:
                self.client = redis.from_url(url, decode_responses=False)
                self.client.ping()
                logger.info("Redis cache connected")
            except Exception as e:
                logger.warning(f"Redis connection failed: {e}. Caching disabled.")
                self.enabled = False
                self.client = None
        else:
            self.client = None
    
    def get(self, key: str) -> Optional[Any]:
        """Get value from cache"""
        if not self.enabled:
            return None
        
        try:
            data = self.client.get(key)
            if data:
                return pickle.loads(data)
            return None
        except Exception as e:
            logger.error(f"Cache get error: {e}")
            return None
    
    def set(self, key: str, value: Any, ttl: int = None) -> bool:
        """Set value in cache"""
        if not self.enabled:
            return False
        
        try:
            ttl = ttl or self.default_ttl
            data = pickle.dumps(value)
            self.client.setex(key, ttl, data)
            return True
        except Exception as e:
            logger.error(f"Cache set error: {e}")
            return False
    
    def delete(self, key: str) -> bool:
        """Delete value from cache"""
        if not self.enabled:
            return False
        
        try:
            self.client.delete(key)
            return True
        except Exception as e:
            logger.error(f"Cache delete error: {e}")
            return False
    
    def delete_pattern(self, pattern: str) -> int:
        """Delete all keys matching pattern"""
        if not self.enabled:
            return 0
        
        try:
            keys = self.client.keys(pattern)
            if keys:
                return self.client.delete(*keys)
            return 0
        except Exception as e:
            logger.error(f"Cache delete pattern error: {e}")
            return 0
    
    def get_embedding(self, text: str) -> Optional[List[float]]:
        """Get cached embedding"""
        key = CacheKeyBuilder.embedding_key(text)
        return self.get(key)
    
    def set_embedding(self, text: str, embedding: List[float]) -> bool:
        """Cache embedding with longer TTL"""
        key = CacheKeyBuilder.embedding_key(text)
        return self.set(key, embedding, ttl=self.embedding_ttl)
    
    def get_query_result(self, query: str, params: dict = None) -> Optional[dict]:
        """Get cached query result"""
        key = CacheKeyBuilder.query_key(query, params)
        return self.get(key)
    
    def set_query_result(self, query: str, result: dict, params: dict = None) -> bool:
        """Cache query result"""
        key = CacheKeyBuilder.query_key(query, params)
        return self.set(key, result)
    
    def invalidate_all(self) -> int:
        """Invalidate all RAG cache entries"""
        return self.delete_pattern("rag:*")
    
    def get_stats(self) -> dict:
        """Get cache statistics"""
        if not self.enabled:
            return {"enabled": False}
        
        try:
            info = self.client.info("memory")
            keys = self.client.dbsize()
            
            return {
                "enabled": True,
                "connected": True,
                "keys": keys,
                "used_memory": info.get("used_memory_human", "unknown"),
                "hit_rate": self._calculate_hit_rate()
            }
        except Exception as e:
            return {"enabled": True, "connected": False, "error": str(e)}
    
    def _calculate_hit_rate(self) -> float:
        """Calculate cache hit rate"""
        try:
            info = self.client.info("stats")
            hits = info.get("keyspace_hits", 0)
            misses = info.get("keyspace_misses", 0)
            total = hits + misses
            return hits / total if total > 0 else 0.0
        except:
            return 0.0


class AsyncRedisCache:
    """Asynchronous Redis cache for high-concurrency scenarios"""
    
    def __init__(self, url: str = REDIS_URL):
        self.enabled = CACHE_CONFIG["enabled"]
        self.default_ttl = CACHE_CONFIG["ttl_seconds"]
        self.url = url
        self.client: Optional[aioredis.Redis] = None
    
    async def connect(self):
        """Initialize async connection"""
        if self.enabled:
            try:
                self.client = await aioredis.from_url(self.url)
                await self.client.ping()
                logger.info("Async Redis cache connected")
            except Exception as e:
                logger.warning(f"Async Redis connection failed: {e}")
                self.enabled = False
    
    async def close(self):
        """Close connection"""
        if self.client:
            await self.client.close()
    
    async def get(self, key: str) -> Optional[Any]:
        """Get value asynchronously"""
        if not self.enabled or not self.client:
            return None
        
        try:
            data = await self.client.get(key)
            if data:
                return pickle.loads(data)
            return None
        except Exception as e:
            logger.error(f"Async cache get error: {e}")
            return None
    
    async def set(self, key: str, value: Any, ttl: int = None) -> bool:
        """Set value asynchronously"""
        if not self.enabled or not self.client:
            return False
        
        try:
            ttl = ttl or self.default_ttl
            data = pickle.dumps(value)
            await self.client.setex(key, ttl, data)
            return True
        except Exception as e:
            logger.error(f"Async cache set error: {e}")
            return False
    
    async def get_many(self, keys: List[str]) -> dict:
        """Get multiple values at once"""
        if not self.enabled or not self.client:
            return {}
        
        try:
            values = await self.client.mget(keys)
            result = {}
            for key, value in zip(keys, values):
                if value:
                    result[key] = pickle.loads(value)
            return result
        except Exception as e:
            logger.error(f"Async cache get_many error: {e}")
            return {}


def cached(ttl: int = None, key_prefix: str = ""):
    """Decorator for caching function results"""
    def decorator(func):
        cache = RedisCache()
        
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not cache.enabled:
                return func(*args, **kwargs)
            
            # Build cache key from function name and arguments
            key_data = {
                "func": func.__name__,
                "args": str(args),
                "kwargs": str(sorted(kwargs.items()))
            }
            key_hash = hashlib.md5(str(key_data).encode()).hexdigest()
            cache_key = f"{key_prefix or 'rag:func'}:{key_hash}"
            
            # Try to get from cache
            cached_result = cache.get(cache_key)
            if cached_result is not None:
                logger.debug(f"Cache hit for {func.__name__}")
                return cached_result
            
            # Execute function and cache result
            result = func(*args, **kwargs)
            cache.set(cache_key, result, ttl)
            logger.debug(f"Cache miss for {func.__name__}, result cached")
            
            return result
        
        return wrapper
    return decorator


def async_cached(ttl: int = None, key_prefix: str = ""):
    """Decorator for caching async function results"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            cache = AsyncRedisCache()
            await cache.connect()
            
            if not cache.enabled:
                return await func(*args, **kwargs)
            
            try:
                # Build cache key
                key_data = {
                    "func": func.__name__,
                    "args": str(args),
                    "kwargs": str(sorted(kwargs.items()))
                }
                key_hash = hashlib.md5(str(key_data).encode()).hexdigest()
                cache_key = f"{key_prefix or 'rag:async'}:{key_hash}"
                
                # Try cache
                cached_result = await cache.get(cache_key)
                if cached_result is not None:
                    return cached_result
                
                # Execute and cache
                result = await func(*args, **kwargs)
                await cache.set(cache_key, result, ttl)
                
                return result
            finally:
                await cache.close()
        
        return wrapper
    return decorator


# Singleton cache instance
_cache_instance: Optional[RedisCache] = None


def get_cache() -> RedisCache:
    """Get global cache instance"""
    global _cache_instance
    if _cache_instance is None:
        _cache_instance = RedisCache()
    return _cache_instance