lethe-core-rust 0.1.1

#!/usr/bin/env python3
"""
Lethe Pool Manager
Manages frozen pools, fingerprints, and validation for regression testing
"""

import os
import json
import hashlib
import subprocess
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from datetime import datetime
import logging

class PoolManager:
    """Manages frozen pools and fingerprints for consistent benchmarking"""
    
    def __init__(self, pools_dir: str = "pools"):
        self.pools_dir = Path(pools_dir)
        self.pools_dir.mkdir(exist_ok=True)
        self.logger = logging.getLogger(__name__)
        
    def compute_pool_fingerprint(self, pool_data: List[Dict]) -> str:
        """Compute deterministic fingerprint for pool data"""
        # Sort pool data by document ID for deterministic hashing
        sorted_data = sorted(pool_data, key=lambda x: x.get('doc_id', ''))
        
        # Create canonical JSON representation
        canonical = json.dumps(sorted_data, sort_keys=True, separators=(',', ':'))
        
        # Compute SHA-256 hash
        fingerprint = hashlib.sha256(canonical.encode('utf-8')).hexdigest()
        return fingerprint[:16]  # Use first 16 chars for readability
    
    def create_frozen_pool(self, documents: List[Dict], metadata: Dict = None) -> str:
        """Create a new frozen pool with fingerprint and manifest"""
        self.logger.info("🧊 Creating frozen pool...")
        
        # Compute fingerprint
        fingerprint = self.compute_pool_fingerprint(documents)
        
        # Create pool data structure
        pool_data = {
            'version': '1.0',
            'fingerprint': fingerprint,
            'created_at': datetime.utcnow().isoformat() + 'Z',
            'document_count': len(documents),
            'metadata': metadata or {},
            'documents': documents
        }
        
        # Save pool data
        pool_file = self.pools_dir / f"frozen_pool_{fingerprint}.jsonl"
        with open(pool_file, 'w') as f:
            for doc in documents:
                f.write(json.dumps(doc) + '\n')
        
        # Create signed manifest
        manifest = {
            'pool_fingerprint': fingerprint,
            'pool_file': str(pool_file),
            'document_count': len(documents),
            'created_at': pool_data['created_at'],
            'tokenizer_hash': self._compute_tokenizer_hash(),
            'schema_version': '1.0',
            'validation_checksum': self._compute_validation_checksum(pool_data),
            'metadata': metadata or {}
        }
        
        manifest_file = self.pools_dir / f"manifest_{fingerprint}.json"
        with open(manifest_file, 'w') as f:
            json.dump(manifest, f, indent=2)
        
        self.logger.info(f"   ✅ Frozen pool created: {fingerprint}")
        self.logger.info(f"   📁 Pool file: {pool_file}")
        self.logger.info(f"   📋 Manifest: {manifest_file}")
        
        return fingerprint
    
    def load_frozen_pool(self, fingerprint: str) -> Optional[List[Dict]]:
        """Load frozen pool by fingerprint"""
        pool_file = self.pools_dir / f"frozen_pool_{fingerprint}.jsonl"
        
        if not pool_file.exists():
            self.logger.error(f"❌ Frozen pool not found: {fingerprint}")
            return None
        
        documents = []
        try:
            with open(pool_file, 'r') as f:
                for line in f:
                    if line.strip():
                        documents.append(json.loads(line))
            
            # Verify fingerprint integrity
            computed_fp = self.compute_pool_fingerprint(documents)
            if computed_fp != fingerprint:
                self.logger.error(f"❌ Pool fingerprint mismatch: {computed_fp} != {fingerprint}")
                return None
            
            self.logger.info(f"   ✅ Loaded frozen pool: {fingerprint} ({len(documents)} docs)")
            return documents
            
        except Exception as e:
            self.logger.error(f"❌ Failed to load frozen pool: {e}")
            return None
    
    def validate_pool_compatibility(self, fingerprint1: str, fingerprint2: str) -> bool:
        """Check if two pools are compatible for comparison"""
        manifest1_file = self.pools_dir / f"manifest_{fingerprint1}.json"
        manifest2_file = self.pools_dir / f"manifest_{fingerprint2}.json"
        
        if not (manifest1_file.exists() and manifest2_file.exists()):
            return False
        
        try:
            with open(manifest1_file, 'r') as f:
                manifest1 = json.load(f)
            with open(manifest2_file, 'r') as f:
                manifest2 = json.load(f)
            
            # Check compatibility criteria
            compatible = (
                manifest1['pool_fingerprint'] == manifest2['pool_fingerprint'] and
                manifest1['tokenizer_hash'] == manifest2['tokenizer_hash'] and
                manifest1['schema_version'] == manifest2['schema_version']
            )
            
            return compatible
            
        except Exception as e:
            self.logger.error(f"❌ Compatibility check failed: {e}")
            return False
    
    def get_latest_pool(self) -> Optional[str]:
        """Get fingerprint of the most recently created pool"""
        manifest_files = list(self.pools_dir.glob("manifest_*.json"))
        
        if not manifest_files:
            return None
        
        # Find most recent manifest
        latest_manifest = max(manifest_files, key=lambda p: p.stat().st_mtime)
        
        try:
            with open(latest_manifest, 'r') as f:
                manifest = json.load(f)
            return manifest['pool_fingerprint']
        except Exception:
            return None
    
    def create_union_pool(self, fingerprints: List[str], name: str = "union") -> Optional[str]:
        """Create a union of multiple frozen pools"""
        self.logger.info(f"🔗 Creating union pool from {len(fingerprints)} pools...")
        
        all_documents = []
        seen_doc_ids = set()
        
        for fp in fingerprints:
            documents = self.load_frozen_pool(fp)
            if documents is None:
                self.logger.error(f"❌ Failed to load pool: {fp}")
                return None
            
            # Deduplicate by doc_id
            for doc in documents:
                doc_id = doc.get('doc_id', '')
                if doc_id not in seen_doc_ids:
                    all_documents.append(doc)
                    seen_doc_ids.add(doc_id)
        
        # Create union pool
        metadata = {
            'type': 'union',
            'source_fingerprints': fingerprints,
            'name': name,
            'deduplication': 'by_doc_id'
        }
        
        union_fingerprint = self.create_frozen_pool(all_documents, metadata)
        self.logger.info(f"   ✅ Union pool created: {union_fingerprint} ({len(all_documents)} docs)")
        
        return union_fingerprint
    
    def _compute_tokenizer_hash(self) -> str:
        """Compute hash of current tokenizer configuration"""
        # This would normally inspect the actual tokenizer configuration
        # For now, return a synthetic hash based on common tokenizer settings
        tokenizer_config = {
            'model': 'bge-small-en-v1.5',
            'max_length': 512,
            'truncation': True,
            'padding': 'max_length',
            'normalize': True
        }
        
        canonical = json.dumps(tokenizer_config, sort_keys=True)
        return hashlib.sha256(canonical.encode()).hexdigest()[:12]
    
    def _compute_validation_checksum(self, pool_data: Dict) -> str:
        """Compute validation checksum for pool integrity"""
        validation_data = {
            'fingerprint': pool_data['fingerprint'],
            'document_count': pool_data['document_count'],
            'created_at': pool_data['created_at']
        }
        
        canonical = json.dumps(validation_data, sort_keys=True)
        return hashlib.sha256(canonical.encode()).hexdigest()[:8]
    
    def list_pools(self) -> List[Dict]:
        """List all available frozen pools"""
        manifest_files = list(self.pools_dir.glob("manifest_*.json"))
        pools = []
        
        for manifest_file in manifest_files:
            try:
                with open(manifest_file, 'r') as f:
                    manifest = json.load(f)
                pools.append(manifest)
            except Exception as e:
                self.logger.warning(f"⚠️  Failed to load manifest {manifest_file}: {e}")
        
        # Sort by creation date
        pools.sort(key=lambda p: p.get('created_at', ''), reverse=True)
        return pools
    
    def cleanup_old_pools(self, keep_count: int = 5) -> int:
        """Clean up old frozen pools, keeping only the most recent ones"""
        pools = self.list_pools()
        
        if len(pools) <= keep_count:
            return 0
        
        to_remove = pools[keep_count:]
        removed_count = 0
        
        for pool in to_remove:
            fingerprint = pool['pool_fingerprint']
            
            # Remove pool file
            pool_file = Path(pool['pool_file'])
            if pool_file.exists():
                pool_file.unlink()
                removed_count += 1
            
            # Remove manifest
            manifest_file = self.pools_dir / f"manifest_{fingerprint}.json"
            if manifest_file.exists():
                manifest_file.unlink()
        
        self.logger.info(f"   🗑️  Removed {removed_count} old pools")
        return removed_count


class PoolValidator:
    """Validates pool integrity and compatibility"""
    
    def __init__(self, pool_manager: PoolManager):
        self.pool_manager = pool_manager
        self.logger = logging.getLogger(__name__)
    
    def validate_pool_integrity(self, fingerprint: str) -> bool:
        """Validate pool file integrity against manifest"""
        self.logger.info(f"🔍 Validating pool integrity: {fingerprint}")
        
        # Load manifest
        manifest_file = self.pool_manager.pools_dir / f"manifest_{fingerprint}.json"
        if not manifest_file.exists():
            self.logger.error("❌ Manifest not found")
            return False
        
        try:
            with open(manifest_file, 'r') as f:
                manifest = json.load(f)
        except Exception as e:
            self.logger.error(f"❌ Failed to load manifest: {e}")
            return False
        
        # Load pool data
        documents = self.pool_manager.load_frozen_pool(fingerprint)
        if documents is None:
            return False
        
        # Validate document count
        if len(documents) != manifest['document_count']:
            self.logger.error(f"❌ Document count mismatch: {len(documents)} != {manifest['document_count']}")
            return False
        
        # Validate fingerprint
        computed_fp = self.pool_manager.compute_pool_fingerprint(documents)
        if computed_fp != fingerprint:
            self.logger.error(f"❌ Fingerprint mismatch: {computed_fp} != {fingerprint}")
            return False
        
        # Validate required fields
        required_doc_fields = ['doc_id', 'content']
        for i, doc in enumerate(documents):
            for field in required_doc_fields:
                if field not in doc:
                    self.logger.error(f"❌ Document {i} missing field: {field}")
                    return False
        
        self.logger.info("   ✅ Pool integrity validated")
        return True
    
    def check_pool_compatibility(self, results_dir: Path) -> Dict[str, any]:
        """Check if benchmark results used compatible pools"""
        compatibility_report = {
            'compatible': True,
            'issues': [],
            'pool_fingerprints': {},
            'recommendations': []
        }
        
        # Scan result directories for pool fingerprints
        result_dirs = [d for d in results_dir.iterdir() if d.is_dir()]
        
        for result_dir in result_dirs:
            result_file = result_dir / "results.json"
            if result_file.exists():
                try:
                    with open(result_file, 'r') as f:
                        results = json.load(f)
                    
                    pool_fp = results.get('pool_fingerprint')
                    if pool_fp:
                        compatibility_report['pool_fingerprints'][result_dir.name] = pool_fp
                        
                except Exception as e:
                    compatibility_report['issues'].append(f"Failed to read {result_file}: {e}")
        
        # Check for mismatches
        fingerprints = list(compatibility_report['pool_fingerprints'].values())
        unique_fingerprints = set(fingerprints)
        
        if len(unique_fingerprints) > 1:
            compatibility_report['compatible'] = False
            compatibility_report['issues'].append("Multiple pool fingerprints detected")
            compatibility_report['recommendations'].append(
                "Use --pool-manifest to force consistent pool usage"
            )
            compatibility_report['recommendations'].append(
                "Regenerate results with frozen pool mode enabled"
            )
        
        return compatibility_report


def main():
    """Main entry point for pool management"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Lethe Pool Manager")
    parser.add_argument("action", choices=[
        'create', 'load', 'list', 'validate', 'union', 'cleanup', 'check-compatibility'
    ])
    parser.add_argument("--fingerprint", help="Pool fingerprint")
    parser.add_argument("--documents", help="JSON file with documents for pool creation")
    parser.add_argument("--name", default="pool", help="Pool name")
    parser.add_argument("--keep-count", type=int, default=5, help="Number of pools to keep during cleanup")
    parser.add_argument("--results-dir", default="benchmark_results", help="Results directory for compatibility check")
    
    args = parser.parse_args()
    
    # Setup logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    
    pool_manager = PoolManager()
    validator = PoolValidator(pool_manager)
    
    if args.action == 'create':
        if not args.documents:
            print("❌ --documents required for pool creation")
            sys.exit(1)
        
        with open(args.documents, 'r') as f:
            documents = json.load(f)
        
        fingerprint = pool_manager.create_frozen_pool(documents, {'name': args.name})
        print(f"✅ Pool created: {fingerprint}")
    
    elif args.action == 'load':
        if not args.fingerprint:
            fingerprint = pool_manager.get_latest_pool()
            if not fingerprint:
                print("❌ No pools found and no fingerprint specified")
                sys.exit(1)
        else:
            fingerprint = args.fingerprint
        
        documents = pool_manager.load_frozen_pool(fingerprint)
        if documents:
            print(f"✅ Loaded pool: {fingerprint} ({len(documents)} documents)")
        else:
            sys.exit(1)
    
    elif args.action == 'list':
        pools = pool_manager.list_pools()
        print(f"Found {len(pools)} frozen pools:")
        for pool in pools:
            print(f"  {pool['pool_fingerprint']}: {pool['document_count']} docs, {pool['created_at']}")
    
    elif args.action == 'validate':
        if not args.fingerprint:
            print("❌ --fingerprint required for validation")
            sys.exit(1)
        
        valid = validator.validate_pool_integrity(args.fingerprint)
        sys.exit(0 if valid else 1)
    
    elif args.action == 'union':
        # For demo, create union of all existing pools
        pools = pool_manager.list_pools()
        if len(pools) < 2:
            print("❌ Need at least 2 pools for union")
            sys.exit(1)
        
        fingerprints = [p['pool_fingerprint'] for p in pools[:3]]  # Use first 3
        union_fp = pool_manager.create_union_pool(fingerprints, args.name)
        if union_fp:
            print(f"✅ Union pool created: {union_fp}")
        else:
            sys.exit(1)
    
    elif args.action == 'cleanup':
        removed = pool_manager.cleanup_old_pools(args.keep_count)
        print(f"✅ Cleaned up {removed} old pools")
    
    elif args.action == 'check-compatibility':
        report = validator.check_pool_compatibility(Path(args.results_dir))
        print(f"Compatible: {report['compatible']}")
        if report['issues']:
            print("Issues:")
            for issue in report['issues']:
                print(f"  - {issue}")
        if report['recommendations']:
            print("Recommendations:")
            for rec in report['recommendations']:
                print(f"  - {rec}")


if __name__ == "__main__":
    main()