""" 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 []