cache/chromadb_cache.py

"""
ChromaDB-based search result cache with vector similarity matching.

This replaces the hash-based cache with a vector database for improved
performance, persistence, and semantic similarity matching.
"""

import os
import json
import time
import uuid
import logging
from typing import Optional, List, Dict, Any
from dataclasses import dataclass, field
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer

logger = logging.getLogger(__name__)

@dataclass
class ChromaCacheEntry:
    """Cache entry for ChromaDB storage"""
    results: List[Dict[str, Any]]
    search_query: str
    search_terms: List[str]
    timestamp: float
    ttl: int  # Time to live in seconds
    hit_count: int = 0
    last_accessed: float = field(default_factory=time.time)
    document_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    
    def is_expired(self) -> bool:
        """Check if cache entry has expired"""
        return time.time() > (self.timestamp + self.ttl)
    
    def is_fresh(self) -> bool:
        """Check if cache entry is still fresh"""
        return not self.is_expired()
    
    def touch(self):
        """Update last accessed time and increment hit count"""
        self.last_accessed = time.time()
        self.hit_count += 1

class ChromaDBSearchCache:
    """ChromaDB-based search result cache with vector similarity matching"""
    
    def __init__(self, 
                 max_size: int = 1000, 
                 default_ttl: int = 3600,
                 cache_db_path: str = "cache_db",
                 cache_results_path: str = "cache_results",
                 embedding_model: str = "all-MiniLM-L6-v2",
                 similarity_threshold: float = 0.7):
        """
        Initialize ChromaDB search cache.
        
        Args:
            max_size: Maximum number of entries in cache
            default_ttl: Default time to live in seconds
            cache_db_path: Path to ChromaDB database directory
            cache_results_path: Path to search results storage directory
            embedding_model: SentenceTransformer model name
            similarity_threshold: Default similarity threshold for matching
        """
        self.max_size = max_size
        self.default_ttl = default_ttl
        self.cache_db_path = cache_db_path
        self.cache_results_path = cache_results_path
        self.similarity_threshold = similarity_threshold
        
        # Initialize embedding model
        self.embedding_model = SentenceTransformer(embedding_model)
        logger.info(f"Loaded SentenceTransformer model: {embedding_model}")
        
        # Initialize ChromaDB client
        self._init_chromadb()
        
        # Statistics tracking
        self.stats = {
            "hits": 0,
            "misses": 0,
            "evictions": 0,
            "expired_evictions": 0,
            "total_entries": 0,
            "vector_searches": 0,
            "exact_matches": 0
        }
        
        # Ensure directories exist
        os.makedirs(self.cache_db_path, exist_ok=True)
        os.makedirs(self.cache_results_path, exist_ok=True)
        
        logger.info(f"ChromaDB cache initialized: max_size={max_size}, ttl={default_ttl}s")
    
    def _init_chromadb(self):
        """Initialize ChromaDB client and collection"""
        try:
            # Initialize ChromaDB client with persistent storage
            self.client = chromadb.PersistentClient(
                path=self.cache_db_path,
                settings=Settings(
                    anonymized_telemetry=False,
                    allow_reset=True
                )
            )
            
            # Get or create collection
            self.collection = self.client.get_or_create_collection(
                name="search_cache_vectors",
                metadata={"description": "Atlas search results cache with vector similarity"}
            )
            
            # Clean up expired entries on startup
            self._cleanup_expired_entries()
            
            logger.info(f"ChromaDB collection initialized: {self.collection.count()} entries")
            
        except Exception as e:
            logger.error(f"Failed to initialize ChromaDB: {e}")
            raise
    
    def _generate_search_text(self, search_terms: List[str]) -> str:
        """Generate search text for embedding from search terms"""
        if not search_terms:
            return ""
        # Join terms with spaces for embedding
        return " ".join(search_terms).lower().strip()
    
    def _cleanup_expired_entries(self):
        """Remove expired entries from ChromaDB and cleanup orphaned files"""
        try:
            current_time = time.time()
            
            # Get all entries
            results = self.collection.get(include=['metadatas', 'documents'])
            
            expired_ids = []
            for i, metadata in enumerate(results.get('metadatas', [])):
                if metadata and 'timestamp' in metadata and 'ttl' in metadata:
                    if current_time > (metadata['timestamp'] + metadata['ttl']):
                        expired_ids.append(results['ids'][i])
            
            if expired_ids:
                # Remove expired entries from ChromaDB
                self.collection.delete(ids=expired_ids)
                
                # Remove associated result files
                for doc_id in expired_ids:
                    result_file = os.path.join(self.cache_results_path, f"{doc_id}.json")
                    if os.path.exists(result_file):
                        os.remove(result_file)
                
                self.stats["expired_evictions"] += len(expired_ids)
                logger.info(f"Cleaned up {len(expired_ids)} expired cache entries")
            
        except Exception as e:
            logger.warning(f"Failed to cleanup expired entries: {e}")
    
    def _evict_lru_entries(self):
        """Evict least recently used entries to make space"""
        try:
            current_count = self.collection.count()
            if current_count < self.max_size:
                return
            
            # Get all entries with metadata
            results = self.collection.get(include=['metadatas'])
            
            # Sort by last_accessed timestamp to find LRU
            entries_with_access = [
                (results['ids'][i], metadata.get('last_accessed', 0))
                for i, metadata in enumerate(results.get('metadatas', []))
                if metadata
            ]
            
            entries_with_access.sort(key=lambda x: x[1])  # Sort by last_accessed
            
            # Calculate how many to evict
            entries_to_evict = current_count - self.max_size + 1
            lru_ids = [entry[0] for entry in entries_with_access[:entries_to_evict]]
            
            if lru_ids:
                # Remove LRU entries
                self.collection.delete(ids=lru_ids)
                
                # Remove associated result files
                for doc_id in lru_ids:
                    result_file = os.path.join(self.cache_results_path, f"{doc_id}.json")
                    if os.path.exists(result_file):
                        os.remove(result_file)
                
                self.stats["evictions"] += len(lru_ids)
                logger.info(f"Evicted {len(lru_ids)} LRU cache entries")
            
        except Exception as e:
            logger.warning(f"Failed to evict LRU entries: {e}")
    
    def _load_search_results(self, document_id: str) -> Optional[List[Dict[str, Any]]]:
        """Load search results from JSON file"""
        try:
            result_file = os.path.join(self.cache_results_path, f"{document_id}.json")
            if os.path.exists(result_file):
                with open(result_file, 'r', encoding='utf-8') as f:
                    return json.load(f)
            return None
        except Exception as e:
            logger.warning(f"Failed to load results for {document_id}: {e}")
            return None
    
    def _save_search_results(self, document_id: str, results: List[Dict[str, Any]]):
        """Save search results to JSON file"""
        try:
            result_file = os.path.join(self.cache_results_path, f"{document_id}.json")
            with open(result_file, 'w', encoding='utf-8') as f:
                json.dump(results, f, indent=2, ensure_ascii=False)
        except Exception as e:
            logger.warning(f"Failed to save results for {document_id}: {e}")
    
    def get(self, search_terms: List[str], 
            use_semantic_matching: bool = True, 
            similarity_threshold: Optional[float] = None) -> Optional[ChromaCacheEntry]:
        """
        Get cached search results using vector similarity matching.
        
        Args:
            search_terms: List of search terms
            use_semantic_matching: Whether to use semantic similarity (always True for ChromaDB)
            similarity_threshold: Similarity threshold for matching (optional)
            
        Returns:
            ChromaCacheEntry if found, None otherwise
        """
        if not search_terms:
            return None
        
        try:
            # Clean up expired entries periodically
            if self.stats["hits"] + self.stats["misses"] % 100 == 0:
                self._cleanup_expired_entries()
            
            # Generate search text for embedding
            search_text = self._generate_search_text(search_terms)
            if not search_text:
                return None
            
            # Use provided threshold or default
            threshold = similarity_threshold or self.similarity_threshold
            
            # Query ChromaDB for similar vectors
            results = self.collection.query(
                query_texts=[search_text],
                n_results=3,  # Get top 3 matches to check TTL
                include=['metadatas', 'documents', 'distances']
            )
            
            self.stats["vector_searches"] += 1
            
            # Check results for valid, non-expired entries
            current_time = time.time()
            for i, (distance, metadata) in enumerate(zip(
                results.get('distances', [[]])[0],
                results.get('metadatas', [[]])[0]
            )):
                if not metadata:
                    continue
                
                # Calculate similarity from distance (ChromaDB uses cosine distance)
                similarity = 1.0 - distance if distance is not None else 0.0
                
                if similarity < threshold:
                    continue
                
                # Check if entry is not expired
                if current_time > (metadata.get('timestamp', 0) + metadata.get('ttl', 0)):
                    continue
                
                # Found valid entry - load results
                document_id = results['ids'][0][i]
                search_results = self._load_search_results(document_id)
                
                if search_results is not None:
                    # Create cache entry
                    search_terms_json = metadata.get('search_terms_json', '[]')
                    try:
                        search_terms = json.loads(search_terms_json)
                    except (json.JSONDecodeError, TypeError):
                        search_terms = []
                    
                    entry = ChromaCacheEntry(
                        results=search_results,
                        search_query=metadata.get('search_query', ''),
                        search_terms=search_terms,
                        timestamp=metadata.get('timestamp', current_time),
                        ttl=metadata.get('ttl', self.default_ttl),
                        hit_count=metadata.get('hit_count', 0),
                        last_accessed=current_time,
                        document_id=document_id
                    )
                    
                    # Update hit count and last_accessed in ChromaDB
                    self.collection.update(
                        ids=[document_id],
                        metadatas=[{
                            **metadata,
                            'hit_count': entry.hit_count + 1,
                            'last_accessed': current_time
                        }]
                    )
                    
                    entry.touch()
                    self.stats["hits"] += 1
                    
                    if similarity > 0.95:
                        self.stats["exact_matches"] += 1
                    
                    logger.info(f"Cache HIT: similarity={similarity:.3f}, age={current_time - entry.timestamp:.0f}s")
                    return entry
            
            # No valid entry found
            self.stats["misses"] += 1
            return None
            
        except Exception as e:
            logger.error(f"Cache get error: {e}")
            self.stats["misses"] += 1
            return None
    
    def put(self, search_terms: List[str], search_query: str, 
            results: List[Dict[str, Any]], ttl: Optional[int] = None):
        """
        Store search results in ChromaDB cache.
        
        Args:
            search_terms: List of search terms
            search_query: Original search query
            results: Search results to cache
            ttl: Time to live in seconds (optional)
        """
        if not search_terms or not results:
            return
        
        try:
            # Use default TTL if not specified
            if ttl is None:
                ttl = self.default_ttl
            
            # Determine TTL based on content type (Phase 3 enhancement)
            query_lower = search_query.lower()
            if any(term in query_lower for term in ["news", "today", "latest", "current", "2024", "2025"]):
                ttl = min(ttl, 900)  # 15 minutes for time-sensitive content
            elif any(term in query_lower for term in ["stock", "price", "rate", "weather"]):
                ttl = min(ttl, 1800)  # 30 minutes for frequently changing data
            
            # Evict old entries if necessary
            self._evict_lru_entries()
            
            # Generate document ID and search text
            document_id = str(uuid.uuid4())
            search_text = self._generate_search_text(search_terms)
            
            current_time = time.time()
            
            # Save search results to file
            self._save_search_results(document_id, results)
            
            # Store in ChromaDB (metadata must be strings, ints, floats, bools, or None)
            self.collection.add(
                documents=[search_text],
                metadatas=[{
                    'search_query': search_query,
                    'search_terms_json': json.dumps(search_terms),  # Convert list to JSON string
                    'timestamp': current_time,
                    'ttl': ttl,
                    'hit_count': 0,
                    'last_accessed': current_time,
                    'result_count': len(results)
                }],
                ids=[document_id]
            )
            
            self.stats["total_entries"] += 1
            
            logger.info(f"Cache STORED: {document_id} (TTL: {ttl}s, Results: {len(results)})")
            
        except Exception as e:
            logger.error(f"Cache put error: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get comprehensive cache statistics"""
        try:
            cache_size = self.collection.count()
            hit_rate = self.stats["hits"] / max(1, self.stats["hits"] + self.stats["misses"]) * 100
            
            # Estimate memory usage
            memory_usage_mb = self._estimate_memory_usage()
            
            return {
                "cache_type": "chromadb_vector",
                "cache_size": cache_size,
                "max_size": self.max_size,
                "hit_rate_percentage": round(hit_rate, 2),
                "total_hits": self.stats["hits"],
                "total_misses": self.stats["misses"],
                "total_evictions": self.stats["evictions"],
                "expired_evictions": self.stats["expired_evictions"],
                "total_entries_created": self.stats["total_entries"],
                "vector_searches": self.stats["vector_searches"],
                "exact_matches": self.stats["exact_matches"],
                "memory_usage_mb": memory_usage_mb,
                "embedding_model": getattr(self.embedding_model, '_model_name', 'all-MiniLM-L6-v2'),
                "similarity_threshold": self.similarity_threshold,
                "persistent_storage": True,
                "database_path": self.cache_db_path,
                "results_path": self.cache_results_path
            }
            
        except Exception as e:
            logger.error(f"Failed to get cache stats: {e}")
            return {"error": str(e)}
    
    def _estimate_memory_usage(self) -> float:
        """Estimate cache memory usage in MB"""
        try:
            # Estimate ChromaDB memory usage
            cache_size = self.collection.count()
            
            # Rough estimates:
            # - Vector storage: 384 dimensions * 4 bytes * count
            # - Metadata: ~500 bytes per entry
            # - File storage not counted (disk-based)
            
            vector_memory = cache_size * 384 * 4  # bytes
            metadata_memory = cache_size * 500  # bytes
            
            total_bytes = vector_memory + metadata_memory
            return round(total_bytes / (1024 * 1024), 2)
            
        except Exception as e:
            logger.warning(f"Failed to estimate memory usage: {e}")
            return 0.0
    
    def clear_expired(self):
        """Manually clear all expired entries"""
        self._cleanup_expired_entries()
        logger.info("Manually cleared expired cache entries")
    
    def clear_all(self):
        """Clear entire cache"""
        try:
            # Delete all documents from collection
            all_results = self.collection.get()
            if all_results.get('ids'):
                self.collection.delete(ids=all_results['ids'])
            
            # Remove all result files
            for filename in os.listdir(self.cache_results_path):
                if filename.endswith('.json'):
                    os.remove(os.path.join(self.cache_results_path, filename))
            
            # Reset stats
            self.stats = {
                "hits": 0,
                "misses": 0,
                "evictions": 0,
                "expired_evictions": 0,
                "total_entries": 0,
                "vector_searches": 0,
                "exact_matches": 0
            }
            
            logger.info("Cache cleared completely")
            
        except Exception as e:
            logger.error(f"Failed to clear cache: {e}")
    
    def get_popular_queries(self, limit: int = 10) -> List[Dict[str, Any]]:
        """Get most popular cached queries by hit count"""
        try:
            results = self.collection.get(include=['metadatas'])
            
            # Sort by hit count
            entries_with_hits = [
                (results['ids'][i], metadata)
                for i, metadata in enumerate(results.get('metadatas', []))
                if metadata and 'hit_count' in metadata
            ]
            
            entries_with_hits.sort(key=lambda x: x[1].get('hit_count', 0), reverse=True)
            
            popular_queries = []
            for i, (doc_id, metadata) in enumerate(entries_with_hits[:limit]):
                try:
                    search_terms = json.loads(metadata.get('search_terms_json', '[]'))
                except (json.JSONDecodeError, TypeError):
                    search_terms = []
                
                popular_queries.append({
                    "rank": i + 1,
                    "document_id": doc_id,
                    "search_query": metadata.get('search_query', ''),
                    "search_terms": search_terms,
                    "hit_count": metadata.get('hit_count', 0),
                    "age_seconds": int(time.time() - metadata.get('timestamp', 0)),
                    "ttl_remaining": max(0, int(metadata.get('ttl', 0) - (time.time() - metadata.get('timestamp', 0))))
                })
            
            return popular_queries
            
        except Exception as e:
            logger.error(f"Failed to get popular queries: {e}")
            return []