""" Collection Manager - Advanced management and optimization for ChromaDB collections This module provides comprehensive collection management including: - Performance optimization - Health monitoring - Maintenance operations - Capacity planning """ import time import json import logging from typing import Dict, List, Any, Optional, Tuple from datetime import datetime, timedelta from collections import defaultdict import numpy as np # MIRA components from core.storage.chroma_client import get_client as get_chroma_client # Configure logging logger = logging.getLogger(__name__) class CollectionManager: """ Advanced collection management system for ChromaDB. Provides optimization, monitoring, and maintenance capabilities to ensure peak performance and reliability of MIRA's consciousness. """ def __init__(self): """Initialize collection manager.""" self.chroma_client = get_chroma_client() # Performance tracking self.performance_history = defaultdict(list) self.optimization_history = [] # Health thresholds self.health_thresholds = { 'max_collection_size': 100000, # Maximum documents per collection 'query_time_warning': 2.0, # Warning if query takes > 2 seconds 'query_time_critical': 5.0, # Critical if query takes > 5 seconds 'embedding_staleness_days': 30, # Re-embed after 30 days 'orphan_threshold': 100, # Cleanup if > 100 orphaned docs 'fragmentation_threshold': 0.2 # 20% fragmentation triggers optimization } logger.info("CollectionManager initialized") def optimize_collections(self) -> Dict[str, Any]: """ Optimize all collections for peak performance. Returns: Optimization results for each collection """ optimization_start = time.time() collections = self.chroma_client.client.list_collections() optimization_results = {} logger.info(f"Starting optimization for {len(collections)} collections") for collection in collections: try: result = self._optimize_collection(collection) optimization_results[collection.name] = result # Track optimization self.optimization_history.append({ 'timestamp': datetime.utcnow().isoformat(), 'collection': collection.name, 'result': result }) except Exception as e: logger.error(f"Error optimizing collection {collection.name}: {e}") optimization_results[collection.name] = { 'status': 'error', 'error': str(e) } total_time = time.time() - optimization_start logger.info(f"Optimization completed in {total_time:.2f}s") return { 'collections': optimization_results, 'total_time': total_time, 'timestamp': datetime.utcnow().isoformat() } def _optimize_collection(self, collection) -> Dict[str, Any]: """ Optimize individual collection. Args: collection: ChromaDB collection object Returns: Optimization results """ start_time = time.time() # Get collection stats stats = self._get_collection_stats(collection) optimizations_applied = [] improvements = {} # Check for orphaned documents if stats['orphaned_docs'] > self.health_thresholds['orphan_threshold']: cleaned = self._cleanup_orphaned_documents(collection, stats) optimizations_applied.append('cleanup_orphaned') improvements['orphaned_cleaned'] = cleaned # Check for outdated embeddings if stats['outdated_embeddings'] > 0: refreshed = self._refresh_embeddings(collection, stats) optimizations_applied.append('refresh_embeddings') improvements['embeddings_refreshed'] = refreshed # Check for metadata inconsistencies if stats['metadata_issues'] > 0: fixed = self._fix_metadata_issues(collection, stats) optimizations_applied.append('fix_metadata') improvements['metadata_fixed'] = fixed # Analyze query performance if stats['avg_query_time'] > self.health_thresholds['query_time_warning']: self._optimize_query_performance(collection) optimizations_applied.append('query_optimization') # Measure improvement end_stats = self._get_collection_stats(collection) performance_improvement = self._calculate_performance_improvement(stats, end_stats) return { 'optimizations_applied': optimizations_applied, 'improvements': improvements, 'performance_improvement': performance_improvement, 'optimization_time': time.time() - start_time, 'status': 'success' } def _get_collection_stats(self, collection) -> Dict[str, Any]: """ Get comprehensive statistics for a collection. Args: collection: ChromaDB collection Returns: Collection statistics """ stats = { 'name': collection.name, 'document_count': collection.count(), 'orphaned_docs': 0, 'outdated_embeddings': 0, 'metadata_issues': 0, 'avg_query_time': 0.0, 'fragmentation': 0.0 } # Sample documents for analysis try: sample_size = min(1000, stats['document_count']) if sample_size > 0: sample = collection.get(limit=sample_size, include=['metadatas', 'embeddings']) # Check for issues stats['orphaned_docs'] = self._count_orphaned_documents(sample) stats['outdated_embeddings'] = self._count_outdated_embeddings(sample) stats['metadata_issues'] = self._count_metadata_issues(sample) # Test query performance stats['avg_query_time'] = self._measure_query_performance(collection) except Exception as e: logger.error(f"Error getting stats for {collection.name}: {e}") return stats def _count_orphaned_documents(self, sample: Dict) -> int: """Count documents that appear to be orphaned.""" orphaned = 0 if not sample or not sample.get('metadatas'): return orphaned for metadata in sample['metadatas']: if metadata: # Check for missing required fields that indicate orphaned state if not metadata.get('timestamp') or not metadata.get('source_data'): orphaned += 1 # Extrapolate to full collection if sample['metadatas']: total_count = len(sample['metadatas']) return int(orphaned * (total_count / len(sample['metadatas']))) return orphaned def _count_outdated_embeddings(self, sample: Dict) -> int: """Count embeddings that are outdated.""" outdated = 0 if not sample or not sample.get('metadatas'): return outdated current_time = datetime.utcnow() staleness_threshold = timedelta(days=self.health_thresholds['embedding_staleness_days']) for metadata in sample['metadatas']: if metadata and metadata.get('timestamp'): try: doc_time = datetime.fromisoformat(metadata['timestamp'].replace('Z', '+00:00')) if current_time - doc_time > staleness_threshold: outdated += 1 except: pass return outdated def _count_metadata_issues(self, sample: Dict) -> int: """Count documents with metadata issues.""" issues = 0 if not sample or not sample.get('metadatas'): return issues for metadata in sample['metadatas']: if metadata: # Check for invalid types (ChromaDB 0.3.29 compatibility) for key, value in metadata.items(): if isinstance(value, (list, dict, bool)): issues += 1 break return issues def _measure_query_performance(self, collection) -> float: """Measure average query performance.""" query_times = [] # Test queries test_queries = [ "test query performance", "neural network implementation", "system optimization" ] for query in test_queries: try: start = time.time() collection.query(query_texts=[query], n_results=10) query_times.append(time.time() - start) except: pass return np.mean(query_times) if query_times else 0.0 def _cleanup_orphaned_documents(self, collection, stats: Dict) -> int: """Clean up orphaned documents.""" cleaned = 0 try: # Get all documents all_docs = collection.get(include=['metadatas']) if not all_docs or not all_docs.get('ids'): return cleaned # Identify orphaned documents orphaned_ids = [] for i, metadata in enumerate(all_docs.get('metadatas', [])): if metadata and (not metadata.get('timestamp') or not metadata.get('source_data')): if i < len(all_docs['ids']): orphaned_ids.append(all_docs['ids'][i]) # Delete in batches batch_size = 100 for i in range(0, len(orphaned_ids), batch_size): batch = orphaned_ids[i:i + batch_size] collection.delete(ids=batch) cleaned += len(batch) logger.info(f"Cleaned {cleaned} orphaned documents from {collection.name}") except Exception as e: logger.error(f"Error cleaning orphaned documents: {e}") return cleaned def _refresh_embeddings(self, collection, stats: Dict) -> int: """Refresh outdated embeddings.""" refreshed = 0 # This would re-generate embeddings for old documents # Placeholder implementation logger.info(f"Would refresh {stats['outdated_embeddings']} embeddings in {collection.name}") return refreshed def _fix_metadata_issues(self, collection, stats: Dict) -> int: """Fix metadata compatibility issues.""" fixed = 0 try: # Get documents with metadata issues all_docs = collection.get(include=['metadatas', 'documents']) if not all_docs or not all_docs.get('ids'): return fixed # Fix metadata issues for i, metadata in enumerate(all_docs.get('metadatas', [])): if metadata: updated = False fixed_metadata = {} for key, value in metadata.items(): if isinstance(value, bool): fixed_metadata[key] = 1 if value else 0 updated = True elif isinstance(value, list): fixed_metadata[key] = ', '.join(str(v) for v in value) updated = True elif isinstance(value, dict): fixed_metadata[key] = json.dumps(value) updated = True else: fixed_metadata[key] = value if updated and i < len(all_docs['ids']): # Update document with fixed metadata collection.update( ids=[all_docs['ids'][i]], metadatas=[fixed_metadata] ) fixed += 1 logger.info(f"Fixed {fixed} metadata issues in {collection.name}") except Exception as e: logger.error(f"Error fixing metadata issues: {e}") return fixed def _optimize_query_performance(self, collection): """Optimize collection for better query performance.""" # This would implement various query optimization strategies # For now, log the need for optimization logger.info(f"Query optimization needed for {collection.name}") def _calculate_performance_improvement(self, before: Dict, after: Dict) -> Dict[str, float]: """Calculate performance improvement metrics.""" improvement = {} # Document count change if before['document_count'] > 0: improvement['size_reduction'] = ( (before['document_count'] - after['document_count']) / before['document_count'] ) # Query time improvement if before['avg_query_time'] > 0: improvement['query_speed_improvement'] = ( (before['avg_query_time'] - after['avg_query_time']) / before['avg_query_time'] ) # Issue resolution improvement['issues_resolved'] = ( before['orphaned_docs'] + before['metadata_issues'] - after['orphaned_docs'] - after['metadata_issues'] ) return improvement def get_collection_health_report(self) -> Dict[str, Any]: """ Generate comprehensive health report for all collections. Returns: Health report with status and recommendations """ report_start = time.time() collections = self.chroma_client.client.list_collections() health_report = { 'timestamp': datetime.utcnow().isoformat(), 'collections': {}, 'overall_health': 'unknown', 'recommendations': [] } health_scores = [] for collection in collections: try: health_data = self._assess_collection_health(collection) health_report['collections'][collection.name] = health_data health_scores.append(health_data['health_score']) except Exception as e: logger.error(f"Error assessing health for {collection.name}: {e}") health_report['collections'][collection.name] = { 'status': 'error', 'error': str(e), 'health_score': 0.0 } # Calculate overall health if health_scores: avg_health = np.mean(health_scores) health_report['overall_health'] = self._get_health_status(avg_health) health_report['average_health_score'] = avg_health # Generate recommendations health_report['recommendations'] = self._generate_health_recommendations(health_report) health_report['report_generation_time'] = time.time() - report_start return health_report def _assess_collection_health(self, collection) -> Dict[str, Any]: """ Assess health of individual collection. Args: collection: ChromaDB collection Returns: Health assessment data """ stats = self._get_collection_stats(collection) # Calculate health score (0-1) health_score = 1.0 issues = [] warnings = [] # Size check if stats['document_count'] > self.health_thresholds['max_collection_size']: health_score -= 0.2 warnings.append(f"Collection size ({stats['document_count']}) exceeds recommended maximum") # Performance check if stats['avg_query_time'] > self.health_thresholds['query_time_critical']: health_score -= 0.3 issues.append(f"Critical query performance: {stats['avg_query_time']:.2f}s average") elif stats['avg_query_time'] > self.health_thresholds['query_time_warning']: health_score -= 0.1 warnings.append(f"Slow query performance: {stats['avg_query_time']:.2f}s average") # Data quality check if stats['orphaned_docs'] > self.health_thresholds['orphan_threshold']: health_score -= 0.1 warnings.append(f"{stats['orphaned_docs']} orphaned documents detected") if stats['metadata_issues'] > 0: health_score -= 0.1 issues.append(f"{stats['metadata_issues']} metadata compatibility issues") if stats['outdated_embeddings'] > stats['document_count'] * 0.1: # > 10% health_score -= 0.1 warnings.append(f"{stats['outdated_embeddings']} outdated embeddings") health_score = max(0.0, health_score) return { 'health_score': health_score, 'status': self._get_health_status(health_score), 'stats': stats, 'issues': issues, 'warnings': warnings, 'last_optimization': self._get_last_optimization_time(collection.name) } def _get_health_status(self, health_score: float) -> str: """Convert health score to status string.""" if health_score >= 0.9: return 'excellent' elif health_score >= 0.7: return 'good' elif health_score >= 0.4: return 'fair' elif health_score >= 0.3: return 'poor' else: return 'critical' def _get_last_optimization_time(self, collection_name: str) -> Optional[str]: """Get last optimization time for collection.""" for optimization in reversed(self.optimization_history): if optimization.get('collection') == collection_name: return optimization.get('timestamp') return None def _generate_health_recommendations(self, health_report: Dict) -> List[str]: """Generate actionable recommendations based on health report.""" recommendations = [] # Check overall health overall_health = health_report.get('overall_health', 'unknown') if overall_health in ['poor', 'critical']: recommendations.append("URGENT: System health is degraded. Run optimization immediately.") # Check individual collections for collection_name, health_data in health_report.get('collections', {}).items(): if health_data.get('status') == 'error': recommendations.append(f"Fix errors in collection '{collection_name}'") continue # Performance recommendations stats = health_data.get('stats', {}) if stats.get('avg_query_time', 0) > self.health_thresholds['query_time_warning']: recommendations.append( f"Optimize query performance for '{collection_name}' - " f"consider indexing or partitioning" ) # Size recommendations if stats.get('document_count', 0) > self.health_thresholds['max_collection_size'] * 0.8: recommendations.append( f"Collection '{collection_name}' approaching size limit - " f"consider archiving old data" ) # Maintenance recommendations if health_data.get('issues'): recommendations.append( f"Address {len(health_data['issues'])} critical issues in '{collection_name}'" ) # Dedup recommendations return list(dict.fromkeys(recommendations)) def schedule_maintenance(self, interval_hours: int = 24) -> Dict[str, Any]: """ Schedule automatic maintenance operations. Args: interval_hours: Hours between maintenance runs Returns: Maintenance schedule information """ # This would integrate with MIRA's daemon system # For now, return schedule information next_run = datetime.utcnow() + timedelta(hours=interval_hours) schedule = { 'interval_hours': interval_hours, 'next_run': next_run.isoformat(), 'operations': [ 'health_check', 'orphan_cleanup', 'metadata_validation', 'performance_optimization' ], 'estimated_duration': '5-10 minutes', 'impact': 'minimal - operations are non-blocking' } logger.info(f"Maintenance scheduled for {next_run.isoformat()}") return schedule def get_performance_trends(self, collection_name: Optional[str] = None, days: int = 7) -> Dict[str, Any]: """ Get performance trends for collections. Args: collection_name: Optional specific collection days: Number of days to analyze Returns: Performance trend data """ # Filter optimization history by date cutoff_date = datetime.utcnow() - timedelta(days=days) trends = { 'period_days': days, 'collections': {} } # Analyze optimization history for record in self.optimization_history: try: record_time = datetime.fromisoformat(record['timestamp']) if record_time < cutoff_date: continue coll_name = record.get('collection', '') if collection_name and coll_name != collection_name: continue if coll_name not in trends['collections']: trends['collections'][coll_name] = { 'optimizations': [], 'improvements': [] } trends['collections'][coll_name]['optimizations'].append({ 'timestamp': record['timestamp'], 'optimizations_applied': record.get('result', {}).get('optimizations_applied', []) }) # Track improvements perf_improvement = record.get('result', {}).get('performance_improvement', {}) if perf_improvement: trends['collections'][coll_name]['improvements'].append(perf_improvement) except Exception as e: logger.error(f"Error processing trend record: {e}") # Calculate summary statistics for coll_name, data in trends['collections'].items(): if data['improvements']: # Average improvements avg_improvements = {} for metric in ['query_speed_improvement', 'size_reduction']: values = [imp.get(metric, 0) for imp in data['improvements'] if metric in imp] if values: avg_improvements[metric] = np.mean(values) data['average_improvements'] = avg_improvements data['optimization_count'] = len(data['optimizations']) return trends