rag/cache.py

"""
Simple cache for RAG queries to speed up repeated questions
"""

from typing import Dict, Any, Optional
import hashlib
import json
from datetime import datetime, timedelta

class QueryCache:
    """Cache for RAG query results"""
    
    def __init__(self, ttl_minutes: int = 60):
        """
        Initialize cache
        
        Args:
            ttl_minutes: Time to live in minutes
        """
        self.cache: Dict[str, Dict[str, Any]] = {}
        self.ttl = timedelta(minutes=ttl_minutes)
    
    def _get_key(self, query: str) -> str:
        """Generate cache key from query"""
        return hashlib.md5(query.lower().strip().encode()).hexdigest()
    
    def get(self, query: str) -> Optional[Dict[str, Any]]:
        """
        Get cached result
        
        Args:
            query: User query
        
        Returns:
            Cached result or None
        """
        key = self._get_key(query)
        
        if key in self.cache:
            entry = self.cache[key]
            # Check if expired
            if datetime.now() - entry['timestamp'] < self.ttl:
                return entry['result']
            else:
                # Remove expired entry
                del self.cache[key]
        
        return None
    
    def set(self, query: str, result: Dict[str, Any]):
        """
        Cache result
        
        Args:
            query: User query
            result: Query result
        """
        key = self._get_key(query)
        self.cache[key] = {
            'result': result,
            'timestamp': datetime.now()
        }
    
    def clear(self):
        """Clear all cache"""
        self.cache.clear()
    
    def size(self) -> int:
        """Get cache size"""
        return len(self.cache)


# Global cache instance
_cache = QueryCache(ttl_minutes=60)

def get_cache() -> QueryCache:
    """Get global cache instance"""
    return _cache