framework.py

"""
Comprehensive LLM Evaluation & Testing Framework

This module provides a complete evaluation framework for LLM-based systems,
supporting multiple metrics, test suites, benchmarking, and A/B testing.
"""
import os
import json
import logging
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple, Callable
from dataclasses import dataclass, asdict, field
from datetime import datetime
from enum import Enum
import time

from dotenv import load_dotenv

load_dotenv()

logger = logging.getLogger(__name__)


class MetricType(Enum):
    """Types of evaluation metrics."""
    RETRIEVAL = "retrieval"
    GENERATION = "generation"
    SEMANTIC = "semantic"
    LLM_JUDGE = "llm_judge"
    CUSTOM = "custom"


class EvaluationMode(Enum):
    """Evaluation execution modes."""
    SINGLE = "single"  # Evaluate single model/config
    AB_TEST = "ab_test"  # Compare two models/configs
    BATCH = "batch"  # Evaluate multiple models/configs


@dataclass
class TestCase:
    """A single test case for evaluation."""
    id: str
    question: str
    expected_answer: Optional[str] = None
    expected_keywords: List[str] = field(default_factory=list)
    expected_sources: List[str] = field(default_factory=list)
    context: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    category: str = "general"
    difficulty: str = "medium"  # easy, medium, hard
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TestCase":
        """Create from dictionary."""
        return cls(**data)


@dataclass
class EvaluationResult:
    """Result of evaluating a single test case."""
    test_case_id: str
    question: str
    predicted_answer: str
    expected_answer: Optional[str]
    metrics: Dict[str, float] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)
    execution_time: float = 0.0
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return asdict(self)


@dataclass
class TestSuite:
    """A collection of test cases."""
    name: str
    description: str
    test_cases: List[TestCase] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    updated_at: str = field(default_factory=lambda: datetime.now().isoformat())
    
    def add_test_case(self, test_case: TestCase):
        """Add a test case to the suite."""
        self.test_cases.append(test_case)
        self.updated_at = datetime.now().isoformat()
    
    def get_test_cases_by_category(self, category: str) -> List[TestCase]:
        """Get test cases filtered by category."""
        return [tc for tc in self.test_cases if tc.category == category]
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "test_cases": [tc.to_dict() for tc in self.test_cases],
            "metadata": self.metadata,
            "created_at": self.created_at,
            "updated_at": self.updated_at
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TestSuite":
        """Create from dictionary."""
        test_cases = [TestCase.from_dict(tc) for tc in data.get("test_cases", [])]
        return cls(
            name=data["name"],
            description=data.get("description", ""),
            test_cases=test_cases,
            metadata=data.get("metadata", {}),
            created_at=data.get("created_at", datetime.now().isoformat()),
            updated_at=data.get("updated_at", datetime.now().isoformat())
        )


class MetricCalculator:
    """Calculates various evaluation metrics."""
    
    @staticmethod
    def calculate_recall_at_k(
        retrieved_items: List[str],
        expected_items: List[str],
        k: int = 5
    ) -> float:
        """Calculate recall@k."""
        if not expected_items:
            return 0.0
        
        top_k = retrieved_items[:k]
        found = sum(1 for item in expected_items if item in top_k)
        return found / len(expected_items)
    
    @staticmethod
    def calculate_precision_at_k(
        retrieved_items: List[str],
        expected_items: List[str],
        k: int = 5
    ) -> float:
        """Calculate precision@k."""
        if not retrieved_items:
            return 0.0
        
        top_k = retrieved_items[:k]
        found = sum(1 for item in top_k if item in expected_items)
        return found / len(top_k)
    
    @staticmethod
    def calculate_hit_rate(
        retrieved_items: List[str],
        expected_items: List[str]
    ) -> float:
        """Calculate hit rate (at least one expected item retrieved)."""
        if not expected_items:
            return 0.0
        
        return 1.0 if any(item in retrieved_items for item in expected_items) else 0.0
    
    @staticmethod
    def calculate_mrr(
        retrieved_items: List[str],
        expected_items: List[str]
    ) -> float:
        """Calculate Mean Reciprocal Rank."""
        if not expected_items:
            return 0.0
        
        for rank, item in enumerate(retrieved_items, 1):
            if item in expected_items:
                return 1.0 / rank
        
        return 0.0
    
    @staticmethod
    def calculate_keyword_overlap(
        predicted_text: str,
        expected_keywords: List[str]
    ) -> float:
        """Calculate keyword overlap score."""
        if not expected_keywords:
            return 0.0
        
        predicted_lower = predicted_text.lower()
        found_keywords = sum(
            1 for kw in expected_keywords
            if kw.lower() in predicted_lower
        )
        return found_keywords / len(expected_keywords)
    
    @staticmethod
    def calculate_exact_match(
        predicted: str,
        expected: str
    ) -> float:
        """Calculate exact match score."""
        return 1.0 if predicted.strip().lower() == expected.strip().lower() else 0.0
    
    @staticmethod
    def calculate_contains_match(
        predicted: str,
        expected: str
    ) -> float:
        """Check if predicted contains expected."""
        return 1.0 if expected.lower() in predicted.lower() else 0.0


class LLMEvaluator:
    """Main evaluation framework for LLM systems."""
    
    def __init__(
        self,
        qa_chain: Any = None,
        vector_store: Any = None,
        output_dir: Optional[Path] = None
    ):
        """
        Initialize the evaluator.
        
        Args:
            qa_chain: QA chain instance (from app.llm.qa_chain)
            vector_store: Vector store instance (from app.retrieval.vector_store)
            output_dir: Directory for saving evaluation results
        """
        self.qa_chain = qa_chain
        self.vector_store = vector_store
        self.output_dir = output_dir or Path("evaluation_results")
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        self.metric_calculator = MetricCalculator()
        self.test_suites: Dict[str, TestSuite] = {}
        self.evaluation_history: List[Dict[str, Any]] = []
    
    def load_test_suite(self, suite_path: Path) -> TestSuite:
        """Load a test suite from JSON file."""
        with open(suite_path, 'r') as f:
            data = json.load(f)
            suite = TestSuite.from_dict(data)
            self.test_suites[suite.name] = suite
            return suite
    
    def save_test_suite(self, suite: TestSuite, suite_path: Optional[Path] = None):
        """Save a test suite to JSON file."""
        if suite_path is None:
            suite_path = self.output_dir / f"{suite.name}.json"
        
        suite_path.parent.mkdir(parents=True, exist_ok=True)
        with open(suite_path, 'w') as f:
            json.dump(suite.to_dict(), f, indent=2)
    
    def evaluate_test_case(
        self,
        test_case: TestCase,
        calculate_metrics: bool = True
    ) -> EvaluationResult:
        """
        Evaluate a single test case.
        
        Args:
            test_case: Test case to evaluate
            calculate_metrics: Whether to calculate metrics
            
        Returns:
            EvaluationResult with metrics
        """
        start_time = time.time()
        
        # Get prediction from QA chain
        if self.qa_chain is None:
            raise ValueError("QA chain not initialized")
        
        try:
            result = self.qa_chain.answer_question(
                test_case.question,
                portfolio_id=test_case.metadata.get("portfolio_id")
            )
            predicted_answer = result.get("answer", "")
            retrieved_sources = [
                src.get("source", "") for src in result.get("sources", [])
            ]
        except Exception as e:
            logger.error(f"Error evaluating test case {test_case.id}: {e}")
            predicted_answer = f"Error: {str(e)}"
            retrieved_sources = []
        
        execution_time = time.time() - start_time
        
        # Calculate metrics
        metrics = {}
        if calculate_metrics:
            # Retrieval metrics
            if test_case.expected_sources:
                metrics["recall_at_5"] = self.metric_calculator.calculate_recall_at_k(
                    retrieved_sources, test_case.expected_sources, k=5
                )
                metrics["recall_at_10"] = self.metric_calculator.calculate_recall_at_k(
                    retrieved_sources, test_case.expected_sources, k=10
                )
                metrics["hit_rate"] = self.metric_calculator.calculate_hit_rate(
                    retrieved_sources, test_case.expected_sources
                )
                metrics["mrr"] = self.metric_calculator.calculate_mrr(
                    retrieved_sources, test_case.expected_sources
                )
            
            # Keyword overlap
            if test_case.expected_keywords:
                metrics["keyword_overlap"] = self.metric_calculator.calculate_keyword_overlap(
                    predicted_answer, test_case.expected_keywords
                )
            
            # Exact/contains match
            if test_case.expected_answer:
                metrics["exact_match"] = self.metric_calculator.calculate_exact_match(
                    predicted_answer, test_case.expected_answer
                )
                metrics["contains_match"] = self.metric_calculator.calculate_contains_match(
                    predicted_answer, test_case.expected_answer
                )
        
        return EvaluationResult(
            test_case_id=test_case.id,
            question=test_case.question,
            predicted_answer=predicted_answer,
            expected_answer=test_case.expected_answer,
            metrics=metrics,
            metadata={
                "retrieved_sources": retrieved_sources,
                "num_sources": len(retrieved_sources),
                "category": test_case.category,
                "difficulty": test_case.difficulty
            },
            execution_time=execution_time
        )
    
    def evaluate_test_suite(
        self,
        suite: TestSuite,
        save_results: bool = True
    ) -> Dict[str, Any]:
        """
        Evaluate an entire test suite.
        
        Args:
            suite: Test suite to evaluate
            save_results: Whether to save results to file
            
        Returns:
            Dictionary with aggregate metrics and per-case results
        """
        logger.info(f"Evaluating test suite: {suite.name} ({len(suite.test_cases)} test cases)")
        
        results = []
        for i, test_case in enumerate(suite.test_cases):
            logger.info(f"Evaluating test case {i+1}/{len(suite.test_cases)}: {test_case.id}")
            result = self.evaluate_test_case(test_case)
            results.append(result)
        
        # Calculate aggregate metrics
        aggregate_metrics = self._calculate_aggregate_metrics(results)
        
        evaluation_summary = {
            "suite_name": suite.name,
            "timestamp": datetime.now().isoformat(),
            "num_test_cases": len(results),
            "aggregate_metrics": aggregate_metrics,
            "results": [r.to_dict() for r in results]
        }
        
        # Save results
        if save_results:
            results_path = self.output_dir / f"evaluation_{suite.name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
            with open(results_path, 'w') as f:
                json.dump(evaluation_summary, f, indent=2)
            logger.info(f"Results saved to {results_path}")
        
        self.evaluation_history.append(evaluation_summary)
        return evaluation_summary
    
    def _calculate_aggregate_metrics(
        self,
        results: List[EvaluationResult]
    ) -> Dict[str, float]:
        """Calculate aggregate metrics across all results."""
        if not results:
            return {}
        
        aggregate = {}
        
        # Get all metric keys
        all_metric_keys = set()
        for result in results:
            all_metric_keys.update(result.metrics.keys())
        
        # Calculate averages for each metric
        for metric_key in all_metric_keys:
            values = [r.metrics.get(metric_key, 0.0) for r in results]
            aggregate[f"avg_{metric_key}"] = sum(values) / len(values) if values else 0.0
        
        # Calculate execution statistics
        execution_times = [r.execution_time for r in results]
        aggregate["avg_execution_time"] = sum(execution_times) / len(execution_times)
        aggregate["total_execution_time"] = sum(execution_times)
        
        return aggregate
    
    def compare_evaluations(
        self,
        evaluation1: Dict[str, Any],
        evaluation2: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Compare two evaluation results (A/B testing).
        
        Args:
            evaluation1: First evaluation results
            evaluation2: Second evaluation results
            
        Returns:
            Comparison results with differences
        """
        comparison = {
            "evaluation1": evaluation1["suite_name"],
            "evaluation2": evaluation2["suite_name"],
            "timestamp": datetime.now().isoformat(),
            "metric_comparisons": {}
        }
        
        metrics1 = evaluation1["aggregate_metrics"]
        metrics2 = evaluation2["aggregate_metrics"]
        
        # Compare each metric
        all_metrics = set(metrics1.keys()) | set(metrics2.keys())
        for metric in all_metrics:
            val1 = metrics1.get(metric, 0.0)
            val2 = metrics2.get(metric, 0.0)
            diff = val2 - val1
            pct_change = (diff / val1 * 100) if val1 != 0 else 0.0
            
            comparison["metric_comparisons"][metric] = {
                "value1": val1,
                "value2": val2,
                "difference": diff,
                "percent_change": pct_change,
                "winner": "evaluation2" if diff > 0 else "evaluation1" if diff < 0 else "tie"
            }
        
        return comparison