src/agents/benchmarking_agent.py

"""
Benchmarking Agent - Tests model performance with real inference
"""

import torch
import time
import asyncio
import numpy as np
from transformers import (
    AutoModelForSequenceClassification, 
    AutoTokenizer, 
    pipeline,
    AutoModelForCausalLM
)
from typing import List, Dict, Any, Optional

from src.models.schemas import (
    TaskType, UserRequirements, BenchmarkResult
)


class BenchmarkingAgent:
    """
    Benchmarks models with real inference tests to measure latency and memory usage.
    
    This agent loads each model, runs warmup inferences, then measures
    performance over multiple iterations.
    """
    
    def __init__(self, sample_data: Dict[str, Any] = None):
        self.sample_data = sample_data or self._get_default_samples()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"   Using device: {self.device}")
    
    async def benchmark_models(self, 
                              model_ids: List[str], 
                              task_type: TaskType,
                              requirements: UserRequirements,
                              max_models: int = 3) -> List[BenchmarkResult]:
        """
        Benchmark top models with quick inference tests.
        
        Args:
            model_ids: List of model IDs to benchmark
            task_type: Type of ML task
            requirements: User requirements
            max_models: Maximum number of models to benchmark
            
        Returns:
            List of BenchmarkResult objects
        """
        results = []
        
        print(f"   Benchmarking up to {max_models} models...")
        
        for i, model_id in enumerate(model_ids[:max_models]):
            print(f"    Testing {i+1}/{min(len(model_ids), max_models)}: {model_id}")
            
            try:
                result = await self._benchmark_single_model(
                    model_id, task_type, requirements
                )
                results.append(result)
                
                if not result.error:
                    print(f"       Latency: {result.latency_ms:.2f}ms, "
                          f"Memory: {result.memory_usage_mb:.2f}MB")
                else:
                    print(f"       Error: {result.error}")
                    
            except Exception as e:
                print(f"       Failed: {e}")
                # FIXED: Added task_type to the error response
                results.append(BenchmarkResult(
                    model_id=model_id,
                    task_type=task_type,  # This was missing!
                    latency_ms=0,
                    memory_usage_mb=0,
                    error=str(e)
                ))
            
            # Small delay between models
            await asyncio.sleep(0.5)
        
        return results
    
    async def _benchmark_single_model(self, 
                                model_id: str,
                                task_type: TaskType,
                                requirements: UserRequirements) -> BenchmarkResult:
        """Benchmark a single model"""
        
        model = None
        tokenizer = None
        nlp_pipeline = None
        
        # Load model and tokenizer
        try:
            print(f"      Loading model...")
            
            if task_type == TaskType.TRANSLATION:
                # For translation models, we need to use pipeline with specific task format
                try:
                    # Try the standard translation pipeline first
                    nlp_pipeline = pipeline(
                        "translation",
                        model=model_id,
                        device=self.device
                    )
                except Exception as e:
                    # If that fails, try with specific language pair format
                    if requirements.translation_reqs:
                        src = requirements.translation_reqs.source_language.value
                        tgt = requirements.translation_reqs.target_language.value
                        task_name = f"translation_{src}_to_{tgt}"
                        try:
                            nlp_pipeline = pipeline(
                                task_name,
                                model=model_id,
                                device=self.device
                            )
                        except:
                            # If both fail, try loading as a general seq2seq model
                            from transformers import AutoModelForSeq2SeqLM
                            tokenizer = AutoTokenizer.from_pretrained(model_id)
                            model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
                
            elif task_type in [TaskType.TEXT_CLASSIFICATION, TaskType.NAMED_ENTITY_RECOGNITION]:
                tokenizer = AutoTokenizer.from_pretrained(model_id)
                model = AutoModelForSequenceClassification.from_pretrained(model_id)
                
            elif task_type == TaskType.TEXT_GENERATION:
                tokenizer = AutoTokenizer.from_pretrained(model_id)
                model = AutoModelForCausalLM.from_pretrained(model_id)
                
            else:
                # Use pipeline for other tasks
                nlp_pipeline = pipeline(
                    task_type.value,
                    model=model_id,
                    device=self.device
                )
                
        except Exception as e:
            return BenchmarkResult(
                model_id=model_id,
                task_type=task_type,
                latency_ms=0,
                memory_usage_mb=0,
                error=f"Failed to load model: {str(e)}"
            )
        
        # Move model to device
        if model:
            model.to(self.device)
            model.eval()
        
        # Get appropriate sample data
        sample = self._get_task_sample(task_type)
        
        # Run warmup (first inference is always slower)
        try:
            await self._run_warmup(model_id, task_type, sample, model, tokenizer, nlp_pipeline, requirements)
        except Exception as e:
            print(f"      Warmup warning: {e}")
        
        # Benchmark inference
        latencies = []
        memory_usage = []
        
        for i in range(5):  # Run 5 iterations for stable measurement
            # Reset memory stats if using CUDA
            if self.device.type == "cuda":
                torch.cuda.reset_peak_memory_stats()
                start_memory = torch.cuda.memory_allocated()
            
            start_time = time.perf_counter()
            
            # Run inference
            try:
                with torch.no_grad():
                    if task_type == TaskType.TRANSLATION and nlp_pipeline:
                        # For translation pipeline
                        result = nlp_pipeline(sample["text"], max_length=128)
                        
                    elif task_type == TaskType.TRANSLATION and model and tokenizer:
                        # For seq2seq model
                        inputs = tokenizer(
                            sample["text"], 
                            return_tensors="pt", 
                            truncation=True, 
                            max_length=128
                        ).to(self.device)
                        outputs = model.generate(**inputs, max_new_tokens=50)
                        
                    elif task_type == TaskType.TEXT_CLASSIFICATION and model and tokenizer:
                        inputs = tokenizer(
                            sample["text"], 
                            return_tensors="pt", 
                            truncation=True, 
                            max_length=128
                        ).to(self.device)
                        outputs = model(**inputs)
                        
                    elif task_type == TaskType.TEXT_GENERATION and model and tokenizer:
                        inputs = tokenizer(
                            sample["text"], 
                            return_tensors="pt", 
                            truncation=True
                        ).to(self.device)
                        outputs = model.generate(**inputs, max_new_tokens=20)
                        
                    elif nlp_pipeline:
                        result = nlp_pipeline(sample["text"])
                        
            except Exception as e:
                return BenchmarkResult(
                    model_id=model_id,
                    task_type=task_type,
                    latency_ms=0,
                    memory_usage_mb=0,
                    error=f"Inference failed: {str(e)}"
                )
            
            end_time = time.perf_counter()
            
            # Measure memory
            if self.device.type == "cuda":
                end_memory = torch.cuda.memory_allocated()
                peak_memory = torch.cuda.max_memory_allocated()
                memory_used = (peak_memory - start_memory) / (1024 ** 2)  # Convert to MB
                memory_usage.append(memory_used)
            
            latency_ms = (end_time - start_time) * 1000
            latencies.append(latency_ms)
            
            # Small delay between runs
            await asyncio.sleep(0.1)
        
        # Clean up
        if model:
            del model
        if tokenizer:
            del tokenizer
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Calculate statistics
        avg_latency = float(np.mean(latencies))
        avg_memory = float(np.mean(memory_usage)) if memory_usage else 0
        
        return BenchmarkResult(
            model_id=model_id,
            task_type=task_type,
            latency_ms=avg_latency,
            memory_usage_mb=avg_memory,
            throughput=1000 / avg_latency if avg_latency > 0 else 0
        )
    
    async def _run_warmup(self, model_id: str, task_type: TaskType, sample: Dict,
                     model=None, tokenizer=None, nlp_pipeline=None, requirements=None):
        """Run warmup inference to initialize model"""
        try:
            with torch.no_grad():
                if task_type == TaskType.TRANSLATION and nlp_pipeline:
                    nlp_pipeline(sample["text"], max_length=50)
                    
                elif task_type == TaskType.TRANSLATION and model and tokenizer:
                    inputs = tokenizer(
                        sample["text"], 
                        return_tensors="pt", 
                        truncation=True
                    ).to(self.device)
                    model.generate(**inputs, max_new_tokens=20)
                    
                elif task_type == TaskType.TEXT_CLASSIFICATION and model and tokenizer:
                    inputs = tokenizer(
                        sample["text"], 
                        return_tensors="pt", 
                        truncation=True
                    ).to(self.device)
                    model(**inputs)
                    
                elif task_type == TaskType.TEXT_GENERATION and model and tokenizer:
                    inputs = tokenizer(
                        sample["text"], 
                        return_tensors="pt", 
                        truncation=True
                    ).to(self.device)
                    model.generate(**inputs, max_new_tokens=10)
                    
                elif nlp_pipeline:
                    nlp_pipeline(sample["text"])
                    
        except Exception as e:
            raise e
    
    def _get_task_sample(self, task_type: TaskType) -> Dict[str, Any]:
        """Get sample data for benchmarking"""
        samples = {
            TaskType.TEXT_CLASSIFICATION: {
                "text": "This is a sample text for classification benchmarking."
            },
            TaskType.TEXT_GENERATION: {
                "text": "Once upon a time in a land far away",
            },
            TaskType.SUMMARIZATION: {
                "text": """Artificial intelligence is transforming industries across the globe. 
                From healthcare to finance, AI systems are being deployed to solve complex problems. 
                Machine learning algorithms can now diagnose diseases, predict market trends, 
                and even create art. The rapid advancement of AI technology brings both opportunities 
                and challenges that society must address."""
            },
            TaskType.QUESTION_ANSWERING: {
                "context": "The Eiffel Tower is located in Paris, France.",
                "question": "Where is the Eiffel Tower?"
            },
            TaskType.TRANSLATION: {
                "text": "Hello, how are you today?"
            },
            TaskType.TEXT_TO_SPEECH: {
                "text": "Hello, this is a test of the text to speech system."
            },
            TaskType.SPEECH_TO_TEXT: {
                "text": "This is a sample audio transcription test."
            },
            TaskType.OCR: {
                "text": "Sample text from an image."
            }
        }
        
        return samples.get(task_type, {"text": "Sample text for benchmarking."})
    
    def _get_default_samples(self) -> Dict[str, Any]:
        """Get default sample data for various tasks"""
        return {
            "text_classification": [
                {"text": "I love this product, it's amazing!", "label": "positive"},
                {"text": "This is the worst experience ever.", "label": "negative"}
            ],
            "summarization": [
                {"text": "Long article about AI advancements..."}
            ],
            "translation": [
                {"text": "Hello world", "source_lang": "en", "target_lang": "fr"}
            ]
        }