findEthics/Atlas
0
1"""2ChromaDB-based search result cache with vector similarity matching.3 4This replaces the hash-based cache with a vector database for improved5performance, persistence, and semantic similarity matching.6"""7 8import os9import json10import time11import uuid12import logging13from typing import Optional, List, Dict, Any14from dataclasses import dataclass, field15import chromadb16from chromadb.config import Settings17from sentence_transformers import SentenceTransformer18 19logger = logging.getLogger(__name__)20 21@dataclass22class ChromaCacheEntry:23 """Cache entry for ChromaDB storage"""24 results: List[Dict[str, Any]]25 search_query: str26 search_terms: List[str]27 timestamp: float28 ttl: int # Time to live in seconds29 hit_count: int = 030 last_accessed: float = field(default_factory=time.time)31 document_id: str = field(default_factory=lambda: str(uuid.uuid4()))32 33 def is_expired(self) -> bool:34 """Check if cache entry has expired"""35 return time.time() > (self.timestamp + self.ttl)36 37 def is_fresh(self) -> bool:38 """Check if cache entry is still fresh"""39 return not self.is_expired()40 41 def touch(self):42 """Update last accessed time and increment hit count"""43 self.last_accessed = time.time()44 self.hit_count += 145 46class ChromaDBSearchCache:47 """ChromaDB-based search result cache with vector similarity matching"""48 49 def __init__(self, 50 max_size: int = 1000, 51 default_ttl: int = 3600,52 cache_db_path: str = "cache_db",53 cache_results_path: str = "cache_results",54 embedding_model: str = "all-MiniLM-L6-v2",55 similarity_threshold: float = 0.7):56 """57 Initialize ChromaDB search cache.58 59 Args:60 max_size: Maximum number of entries in cache61 default_ttl: Default time to live in seconds62 cache_db_path: Path to ChromaDB database directory63 cache_results_path: Path to search results storage directory64 embedding_model: SentenceTransformer model name65 similarity_threshold: Default similarity threshold for matching66 """67 self.max_size = max_size68 self.default_ttl = default_ttl69 self.cache_db_path = cache_db_path70 self.cache_results_path = cache_results_path71 self.similarity_threshold = similarity_threshold72 73 # Initialize embedding model74 self.embedding_model = SentenceTransformer(embedding_model)75 logger.info(f"Loaded SentenceTransformer model: {embedding_model}")76 77 # Initialize ChromaDB client78 self._init_chromadb()79 80 # Statistics tracking81 self.stats = {82 "hits": 0,83 "misses": 0,84 "evictions": 0,85 "expired_evictions": 0,86 "total_entries": 0,87 "vector_searches": 0,88 "exact_matches": 089 }90 91 # Ensure directories exist92 os.makedirs(self.cache_db_path, exist_ok=True)93 os.makedirs(self.cache_results_path, exist_ok=True)94 95 logger.info(f"ChromaDB cache initialized: max_size={max_size}, ttl={default_ttl}s")96 97 def _init_chromadb(self):98 """Initialize ChromaDB client and collection"""99 try:100 # Initialize ChromaDB client with persistent storage101 self.client = chromadb.PersistentClient(102 path=self.cache_db_path,103 settings=Settings(104 anonymized_telemetry=False,105 allow_reset=True106 )107 )108 109 # Get or create collection110 self.collection = self.client.get_or_create_collection(111 name="search_cache_vectors",112 metadata={"description": "Atlas search results cache with vector similarity"}113 )114 115 # Clean up expired entries on startup116 self._cleanup_expired_entries()117 118 logger.info(f"ChromaDB collection initialized: {self.collection.count()} entries")119 120 except Exception as e:121 logger.error(f"Failed to initialize ChromaDB: {e}")122 raise123 124 def _generate_search_text(self, search_terms: List[str]) -> str:125 """Generate search text for embedding from search terms"""126 if not search_terms:127 return ""128 # Join terms with spaces for embedding129 return " ".join(search_terms).lower().strip()130 131 def _cleanup_expired_entries(self):132 """Remove expired entries from ChromaDB and cleanup orphaned files"""133 try:134 current_time = time.time()135 136 # Get all entries137 results = self.collection.get(include=['metadatas', 'documents'])138 139 expired_ids = []140 for i, metadata in enumerate(results.get('metadatas', [])):141 if metadata and 'timestamp' in metadata and 'ttl' in metadata:142 if current_time > (metadata['timestamp'] + metadata['ttl']):143 expired_ids.append(results['ids'][i])144 145 if expired_ids:146 # Remove expired entries from ChromaDB147 self.collection.delete(ids=expired_ids)148 149 # Remove associated result files150 for doc_id in expired_ids:151 result_file = os.path.join(self.cache_results_path, f"{doc_id}.json")152 if os.path.exists(result_file):153 os.remove(result_file)154 155 self.stats["expired_evictions"] += len(expired_ids)156 logger.info(f"Cleaned up {len(expired_ids)} expired cache entries")157 158 except Exception as e:159 logger.warning(f"Failed to cleanup expired entries: {e}")160 161 def _evict_lru_entries(self):162 """Evict least recently used entries to make space"""163 try:164 current_count = self.collection.count()165 if current_count < self.max_size:166 return167 168 # Get all entries with metadata169 results = self.collection.get(include=['metadatas'])170 171 # Sort by last_accessed timestamp to find LRU172 entries_with_access = [173 (results['ids'][i], metadata.get('last_accessed', 0))174 for i, metadata in enumerate(results.get('metadatas', []))175 if metadata176 ]177 178 entries_with_access.sort(key=lambda x: x[1]) # Sort by last_accessed179 180 # Calculate how many to evict181 entries_to_evict = current_count - self.max_size + 1182 lru_ids = [entry[0] for entry in entries_with_access[:entries_to_evict]]183 184 if lru_ids:185 # Remove LRU entries186 self.collection.delete(ids=lru_ids)187 188 # Remove associated result files189 for doc_id in lru_ids:190 result_file = os.path.join(self.cache_results_path, f"{doc_id}.json")191 if os.path.exists(result_file):192 os.remove(result_file)193 194 self.stats["evictions"] += len(lru_ids)195 logger.info(f"Evicted {len(lru_ids)} LRU cache entries")196 197 except Exception as e:198 logger.warning(f"Failed to evict LRU entries: {e}")199 200 def _load_search_results(self, document_id: str) -> Optional[List[Dict[str, Any]]]:201 """Load search results from JSON file"""202 try:203 result_file = os.path.join(self.cache_results_path, f"{document_id}.json")204 if os.path.exists(result_file):205 with open(result_file, 'r', encoding='utf-8') as f:206 return json.load(f)207 return None208 except Exception as e:209 logger.warning(f"Failed to load results for {document_id}: {e}")210 return None211 212 def _save_search_results(self, document_id: str, results: List[Dict[str, Any]]):213 """Save search results to JSON file"""214 try:215 result_file = os.path.join(self.cache_results_path, f"{document_id}.json")216 with open(result_file, 'w', encoding='utf-8') as f:217 json.dump(results, f, indent=2, ensure_ascii=False)218 except Exception as e:219 logger.warning(f"Failed to save results for {document_id}: {e}")220 221 def get(self, search_terms: List[str], 222 use_semantic_matching: bool = True, 223 similarity_threshold: Optional[float] = None) -> Optional[ChromaCacheEntry]:224 """225 Get cached search results using vector similarity matching.226 227 Args:228 search_terms: List of search terms229 use_semantic_matching: Whether to use semantic similarity (always True for ChromaDB)230 similarity_threshold: Similarity threshold for matching (optional)231 232 Returns:233 ChromaCacheEntry if found, None otherwise234 """235 if not search_terms:236 return None237 238 try:239 # Clean up expired entries periodically240 if self.stats["hits"] + self.stats["misses"] % 100 == 0:241 self._cleanup_expired_entries()242 243 # Generate search text for embedding244 search_text = self._generate_search_text(search_terms)245 if not search_text:246 return None247 248 # Use provided threshold or default249 threshold = similarity_threshold or self.similarity_threshold250 251 # Query ChromaDB for similar vectors252 results = self.collection.query(253 query_texts=[search_text],254 n_results=3, # Get top 3 matches to check TTL255 include=['metadatas', 'documents', 'distances']256 )257 258 self.stats["vector_searches"] += 1259 260 # Check results for valid, non-expired entries261 current_time = time.time()262 for i, (distance, metadata) in enumerate(zip(263 results.get('distances', [[]])[0],264 results.get('metadatas', [[]])[0]265 )):266 if not metadata:267 continue268 269 # Calculate similarity from distance (ChromaDB uses cosine distance)270 similarity = 1.0 - distance if distance is not None else 0.0271 272 if similarity < threshold:273 continue274 275 # Check if entry is not expired276 if current_time > (metadata.get('timestamp', 0) + metadata.get('ttl', 0)):277 continue278 279 # Found valid entry - load results280 document_id = results['ids'][0][i]281 search_results = self._load_search_results(document_id)282 283 if search_results is not None:284 # Create cache entry285 search_terms_json = metadata.get('search_terms_json', '[]')286 try:287 search_terms = json.loads(search_terms_json)288 except (json.JSONDecodeError, TypeError):289 search_terms = []290 291 entry = ChromaCacheEntry(292 results=search_results,293 search_query=metadata.get('search_query', ''),294 search_terms=search_terms,295 timestamp=metadata.get('timestamp', current_time),296 ttl=metadata.get('ttl', self.default_ttl),297 hit_count=metadata.get('hit_count', 0),298 last_accessed=current_time,299 document_id=document_id300 )301 302 # Update hit count and last_accessed in ChromaDB303 self.collection.update(304 ids=[document_id],305 metadatas=[{306 **metadata,307 'hit_count': entry.hit_count + 1,308 'last_accessed': current_time309 }]310 )311 312 entry.touch()313 self.stats["hits"] += 1314 315 if similarity > 0.95:316 self.stats["exact_matches"] += 1317 318 logger.info(f"Cache HIT: similarity={similarity:.3f}, age={current_time - entry.timestamp:.0f}s")319 return entry320 321 # No valid entry found322 self.stats["misses"] += 1323 return None324 325 except Exception as e:326 logger.error(f"Cache get error: {e}")327 self.stats["misses"] += 1328 return None329 330 def put(self, search_terms: List[str], search_query: str, 331 results: List[Dict[str, Any]], ttl: Optional[int] = None):332 """333 Store search results in ChromaDB cache.334 335 Args:336 search_terms: List of search terms337 search_query: Original search query338 results: Search results to cache339 ttl: Time to live in seconds (optional)340 """341 if not search_terms or not results:342 return343 344 try:345 # Use default TTL if not specified346 if ttl is None:347 ttl = self.default_ttl348 349 # Determine TTL based on content type (Phase 3 enhancement)350 query_lower = search_query.lower()351 if any(term in query_lower for term in ["news", "today", "latest", "current", "2024", "2025"]):352 ttl = min(ttl, 900) # 15 minutes for time-sensitive content353 elif any(term in query_lower for term in ["stock", "price", "rate", "weather"]):354 ttl = min(ttl, 1800) # 30 minutes for frequently changing data355 356 # Evict old entries if necessary357 self._evict_lru_entries()358 359 # Generate document ID and search text360 document_id = str(uuid.uuid4())361 search_text = self._generate_search_text(search_terms)362 363 current_time = time.time()364 365 # Save search results to file366 self._save_search_results(document_id, results)367 368 # Store in ChromaDB (metadata must be strings, ints, floats, bools, or None)369 self.collection.add(370 documents=[search_text],371 metadatas=[{372 'search_query': search_query,373 'search_terms_json': json.dumps(search_terms), # Convert list to JSON string374 'timestamp': current_time,375 'ttl': ttl,376 'hit_count': 0,377 'last_accessed': current_time,378 'result_count': len(results)379 }],380 ids=[document_id]381 )382 383 self.stats["total_entries"] += 1384 385 logger.info(f"Cache STORED: {document_id} (TTL: {ttl}s, Results: {len(results)})")386 387 except Exception as e:388 logger.error(f"Cache put error: {e}")389 390 def get_stats(self) -> Dict[str, Any]:391 """Get comprehensive cache statistics"""392 try:393 cache_size = self.collection.count()394 hit_rate = self.stats["hits"] / max(1, self.stats["hits"] + self.stats["misses"]) * 100395 396 # Estimate memory usage397 memory_usage_mb = self._estimate_memory_usage()398 399 return {400 "cache_type": "chromadb_vector",401 "cache_size": cache_size,402 "max_size": self.max_size,403 "hit_rate_percentage": round(hit_rate, 2),404 "total_hits": self.stats["hits"],405 "total_misses": self.stats["misses"],406 "total_evictions": self.stats["evictions"],407 "expired_evictions": self.stats["expired_evictions"],408 "total_entries_created": self.stats["total_entries"],409 "vector_searches": self.stats["vector_searches"],410 "exact_matches": self.stats["exact_matches"],411 "memory_usage_mb": memory_usage_mb,412 "embedding_model": getattr(self.embedding_model, '_model_name', 'all-MiniLM-L6-v2'),413 "similarity_threshold": self.similarity_threshold,414 "persistent_storage": True,415 "database_path": self.cache_db_path,416 "results_path": self.cache_results_path417 }418 419 except Exception as e:420 logger.error(f"Failed to get cache stats: {e}")421 return {"error": str(e)}422 423 def _estimate_memory_usage(self) -> float:424 """Estimate cache memory usage in MB"""425 try:426 # Estimate ChromaDB memory usage427 cache_size = self.collection.count()428 429 # Rough estimates:430 # - Vector storage: 384 dimensions * 4 bytes * count431 # - Metadata: ~500 bytes per entry432 # - File storage not counted (disk-based)433 434 vector_memory = cache_size * 384 * 4 # bytes435 metadata_memory = cache_size * 500 # bytes436 437 total_bytes = vector_memory + metadata_memory438 return round(total_bytes / (1024 * 1024), 2)439 440 except Exception as e:441 logger.warning(f"Failed to estimate memory usage: {e}")442 return 0.0443 444 def clear_expired(self):445 """Manually clear all expired entries"""446 self._cleanup_expired_entries()447 logger.info("Manually cleared expired cache entries")448 449 def clear_all(self):450 """Clear entire cache"""451 try:452 # Delete all documents from collection453 all_results = self.collection.get()454 if all_results.get('ids'):455 self.collection.delete(ids=all_results['ids'])456 457 # Remove all result files458 for filename in os.listdir(self.cache_results_path):459 if filename.endswith('.json'):460 os.remove(os.path.join(self.cache_results_path, filename))461 462 # Reset stats463 self.stats = {464 "hits": 0,465 "misses": 0,466 "evictions": 0,467 "expired_evictions": 0,468 "total_entries": 0,469 "vector_searches": 0,470 "exact_matches": 0471 }472 473 logger.info("Cache cleared completely")474 475 except Exception as e:476 logger.error(f"Failed to clear cache: {e}")477 478 def get_popular_queries(self, limit: int = 10) -> List[Dict[str, Any]]:479 """Get most popular cached queries by hit count"""480 try:481 results = self.collection.get(include=['metadatas'])482 483 # Sort by hit count484 entries_with_hits = [485 (results['ids'][i], metadata)486 for i, metadata in enumerate(results.get('metadatas', []))487 if metadata and 'hit_count' in metadata488 ]489 490 entries_with_hits.sort(key=lambda x: x[1].get('hit_count', 0), reverse=True)491 492 popular_queries = []493 for i, (doc_id, metadata) in enumerate(entries_with_hits[:limit]):494 try:495 search_terms = json.loads(metadata.get('search_terms_json', '[]'))496 except (json.JSONDecodeError, TypeError):497 search_terms = []498 499 popular_queries.append({500 "rank": i + 1,501 "document_id": doc_id,502 "search_query": metadata.get('search_query', ''),503 "search_terms": search_terms,504 "hit_count": metadata.get('hit_count', 0),505 "age_seconds": int(time.time() - metadata.get('timestamp', 0)),506 "ttl_remaining": max(0, int(metadata.get('ttl', 0) - (time.time() - metadata.get('timestamp', 0))))507 })508 509 return popular_queries510 511 except Exception as e:512 logger.error(f"Failed to get popular queries: {e}")513 return []