shared_utils/generation/prompt_optimizer.py

"""
A/B Testing Framework for Prompt Optimization.

This module provides systematic prompt optimization through A/B testing,
performance analysis, and automated prompt variation generation.
"""

import json
import time
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, asdict
from enum import Enum
from pathlib import Path
import numpy as np
from collections import defaultdict
import logging

from .prompt_templates import QueryType, PromptTemplate, TechnicalPromptTemplates


class OptimizationMetric(Enum):
    """Metrics for evaluating prompt performance."""
    RESPONSE_TIME = "response_time"
    CONFIDENCE_SCORE = "confidence_score"
    CITATION_COUNT = "citation_count"
    ANSWER_LENGTH = "answer_length"
    TECHNICAL_ACCURACY = "technical_accuracy"
    USER_SATISFACTION = "user_satisfaction"


@dataclass
class PromptVariation:
    """Represents a prompt variation for A/B testing."""
    variation_id: str
    name: str
    description: str
    template: PromptTemplate
    query_type: QueryType
    created_at: float
    metadata: Dict[str, Any]


@dataclass
class TestResult:
    """Represents a single test result."""
    variation_id: str
    query: str
    query_type: QueryType
    response_time: float
    confidence_score: float
    citation_count: int
    answer_length: int
    technical_accuracy: Optional[float] = None
    user_satisfaction: Optional[float] = None
    timestamp: float = None
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = time.time()
        if self.metadata is None:
            self.metadata = {}


@dataclass
class ComparisonResult:
    """Results of A/B test comparison."""
    variation_a: str
    variation_b: str
    metric: OptimizationMetric
    a_mean: float
    b_mean: float
    improvement_percent: float
    p_value: float
    confidence_interval: Tuple[float, float]
    is_significant: bool
    sample_size: int
    recommendation: str


class PromptOptimizer:
    """
    A/B testing framework for systematic prompt optimization.
    
    Features:
    - Automated prompt variation generation
    - Performance metric tracking
    - Statistical significance testing
    - Recommendation engine
    - Persistence and experiment tracking
    """
    
    def __init__(self, experiment_dir: str = "experiments"):
        """
        Initialize the prompt optimizer.
        
        Args:
            experiment_dir: Directory to store experiment data
        """
        self.experiment_dir = Path(experiment_dir)
        self.experiment_dir.mkdir(exist_ok=True)
        
        self.variations: Dict[str, PromptVariation] = {}
        self.test_results: List[TestResult] = []
        self.active_experiments: Dict[str, List[str]] = {}
        
        # Load existing experiments
        self._load_experiments()
        
        # Setup logging
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)
    
    def create_variation(
        self,
        base_template: PromptTemplate,
        query_type: QueryType,
        variation_name: str,
        modifications: Dict[str, str],
        description: str = ""
    ) -> str:
        """
        Create a new prompt variation.
        
        Args:
            base_template: Base template to modify
            query_type: Type of query this variation is for
            variation_name: Human-readable name
            modifications: Dict of template field modifications
            description: Description of the variation
            
        Returns:
            Variation ID
        """
        variation_id = f"{query_type.value}_{variation_name}_{int(time.time())}"
        
        # Create modified template
        modified_template = PromptTemplate(
            system_prompt=modifications.get("system_prompt", base_template.system_prompt),
            context_format=modifications.get("context_format", base_template.context_format),
            query_format=modifications.get("query_format", base_template.query_format),
            answer_guidelines=modifications.get("answer_guidelines", base_template.answer_guidelines)
        )
        
        variation = PromptVariation(
            variation_id=variation_id,
            name=variation_name,
            description=description,
            template=modified_template,
            query_type=query_type,
            created_at=time.time(),
            metadata=modifications
        )
        
        self.variations[variation_id] = variation
        self._save_variation(variation)
        
        self.logger.info(f"Created variation: {variation_id}")
        return variation_id
    
    def create_temperature_variations(
        self,
        base_query_type: QueryType,
        temperatures: List[float] = [0.3, 0.5, 0.7, 0.9]
    ) -> List[str]:
        """
        Create variations with different temperature settings.
        
        Args:
            base_query_type: Query type to create variations for
            temperatures: List of temperature values to test
            
        Returns:
            List of variation IDs
        """
        base_template = TechnicalPromptTemplates.get_template_for_query("")
        if base_query_type != QueryType.GENERAL:
            template_map = {
                QueryType.DEFINITION: TechnicalPromptTemplates.get_definition_template,
                QueryType.IMPLEMENTATION: TechnicalPromptTemplates.get_implementation_template,
                QueryType.COMPARISON: TechnicalPromptTemplates.get_comparison_template,
                QueryType.SPECIFICATION: TechnicalPromptTemplates.get_specification_template,
                QueryType.CODE_EXAMPLE: TechnicalPromptTemplates.get_code_example_template,
                QueryType.HARDWARE_CONSTRAINT: TechnicalPromptTemplates.get_hardware_constraint_template,
                QueryType.TROUBLESHOOTING: TechnicalPromptTemplates.get_troubleshooting_template,
            }
            base_template = template_map[base_query_type]()
        
        variation_ids = []
        for temp in temperatures:
            temp_modification = {
                "system_prompt": base_template.system_prompt + f"\n\nGenerate responses with temperature={temp} (creativity level).",
                "answer_guidelines": base_template.answer_guidelines + f"\n\nAdjust response creativity to temperature={temp}."
            }
            
            variation_id = self.create_variation(
                base_template=base_template,
                query_type=base_query_type,
                variation_name=f"temp_{temp}",
                modifications=temp_modification,
                description=f"Temperature variation with {temp} creativity level"
            )
            variation_ids.append(variation_id)
        
        return variation_ids
    
    def create_length_variations(
        self,
        base_query_type: QueryType,
        length_styles: List[str] = ["concise", "detailed", "comprehensive"]
    ) -> List[str]:
        """
        Create variations with different response length preferences.
        
        Args:
            base_query_type: Query type to create variations for
            length_styles: List of length styles to test
            
        Returns:
            List of variation IDs
        """
        base_template = TechnicalPromptTemplates.get_template_for_query("")
        if base_query_type != QueryType.GENERAL:
            template_map = {
                QueryType.DEFINITION: TechnicalPromptTemplates.get_definition_template,
                QueryType.IMPLEMENTATION: TechnicalPromptTemplates.get_implementation_template,
                QueryType.COMPARISON: TechnicalPromptTemplates.get_comparison_template,
                QueryType.SPECIFICATION: TechnicalPromptTemplates.get_specification_template,
                QueryType.CODE_EXAMPLE: TechnicalPromptTemplates.get_code_example_template,
                QueryType.HARDWARE_CONSTRAINT: TechnicalPromptTemplates.get_hardware_constraint_template,
                QueryType.TROUBLESHOOTING: TechnicalPromptTemplates.get_troubleshooting_template,
            }
            base_template = template_map[base_query_type]()
        
        length_prompts = {
            "concise": "Be concise and focus on essential information only. Aim for 2-3 sentences per point.",
            "detailed": "Provide detailed explanations with examples. Aim for comprehensive coverage.",
            "comprehensive": "Provide exhaustive detail with multiple examples, edge cases, and related concepts."
        }
        
        variation_ids = []
        for style in length_styles:
            length_modification = {
                "answer_guidelines": base_template.answer_guidelines + f"\n\nResponse style: {length_prompts[style]}"
            }
            
            variation_id = self.create_variation(
                base_template=base_template,
                query_type=base_query_type,
                variation_name=f"length_{style}",
                modifications=length_modification,
                description=f"Length variation with {style} response style"
            )
            variation_ids.append(variation_id)
        
        return variation_ids
    
    def create_citation_variations(
        self,
        base_query_type: QueryType,
        citation_styles: List[str] = ["minimal", "standard", "extensive"]
    ) -> List[str]:
        """
        Create variations with different citation requirements.
        
        Args:
            base_query_type: Query type to create variations for
            citation_styles: List of citation styles to test
            
        Returns:
            List of variation IDs
        """
        base_template = TechnicalPromptTemplates.get_template_for_query("")
        if base_query_type != QueryType.GENERAL:
            template_map = {
                QueryType.DEFINITION: TechnicalPromptTemplates.get_definition_template,
                QueryType.IMPLEMENTATION: TechnicalPromptTemplates.get_implementation_template,
                QueryType.COMPARISON: TechnicalPromptTemplates.get_comparison_template,
                QueryType.SPECIFICATION: TechnicalPromptTemplates.get_specification_template,
                QueryType.CODE_EXAMPLE: TechnicalPromptTemplates.get_code_example_template,
                QueryType.HARDWARE_CONSTRAINT: TechnicalPromptTemplates.get_hardware_constraint_template,
                QueryType.TROUBLESHOOTING: TechnicalPromptTemplates.get_troubleshooting_template,
            }
            base_template = template_map[base_query_type]()
        
        citation_prompts = {
            "minimal": "Use [chunk_X] citations only for direct quotes or specific claims.",
            "standard": "Include [chunk_X] citations for each major point or claim.",
            "extensive": "Provide [chunk_X] citations for every statement. Use multiple citations per point where relevant."
        }
        
        variation_ids = []
        for style in citation_styles:
            citation_modification = {
                "answer_guidelines": base_template.answer_guidelines + f"\n\nCitation style: {citation_prompts[style]}"
            }
            
            variation_id = self.create_variation(
                base_template=base_template,
                query_type=base_query_type,
                variation_name=f"citation_{style}",
                modifications=citation_modification,
                description=f"Citation variation with {style} citation requirements"
            )
            variation_ids.append(variation_id)
        
        return variation_ids
    
    def setup_experiment(
        self,
        experiment_name: str,
        variation_ids: List[str],
        test_queries: List[str]
    ) -> str:
        """
        Set up a new A/B test experiment.
        
        Args:
            experiment_name: Name of the experiment
            variation_ids: List of variation IDs to test
            test_queries: List of test queries
            
        Returns:
            Experiment ID
        """
        experiment_id = f"exp_{experiment_name}_{int(time.time())}"
        
        experiment_config = {
            "experiment_id": experiment_id,
            "name": experiment_name,
            "variation_ids": variation_ids,
            "test_queries": test_queries,
            "created_at": time.time(),
            "status": "active"
        }
        
        self.active_experiments[experiment_id] = variation_ids
        
        # Save experiment config
        experiment_file = self.experiment_dir / f"{experiment_id}.json"
        with open(experiment_file, 'w') as f:
            json.dump(experiment_config, f, indent=2)
        
        self.logger.info(f"Created experiment: {experiment_id}")
        return experiment_id
    
    def record_test_result(
        self,
        variation_id: str,
        query: str,
        query_type: QueryType,
        response_time: float,
        confidence_score: float,
        citation_count: int,
        answer_length: int,
        technical_accuracy: Optional[float] = None,
        user_satisfaction: Optional[float] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """
        Record a test result for analysis.
        
        Args:
            variation_id: ID of the variation tested
            query: The query that was tested
            query_type: Type of the query
            response_time: Response time in seconds
            confidence_score: Confidence score (0-1)
            citation_count: Number of citations in response
            answer_length: Length of answer in characters
            technical_accuracy: Optional technical accuracy score (0-1)
            user_satisfaction: Optional user satisfaction score (0-1)
            metadata: Optional additional metadata
        """
        result = TestResult(
            variation_id=variation_id,
            query=query,
            query_type=query_type,
            response_time=response_time,
            confidence_score=confidence_score,
            citation_count=citation_count,
            answer_length=answer_length,
            technical_accuracy=technical_accuracy,
            user_satisfaction=user_satisfaction,
            metadata=metadata or {}
        )
        
        self.test_results.append(result)
        self._save_test_result(result)
        
        self.logger.info(f"Recorded test result for variation: {variation_id}")
    
    def analyze_variations(
        self,
        variation_a: str,
        variation_b: str,
        metric: OptimizationMetric,
        min_samples: int = 10
    ) -> ComparisonResult:
        """
        Analyze performance difference between two variations.
        
        Args:
            variation_a: First variation ID
            variation_b: Second variation ID
            metric: Metric to compare
            min_samples: Minimum samples required for analysis
            
        Returns:
            Comparison result with statistical analysis
        """
        # Filter results for each variation
        results_a = [r for r in self.test_results if r.variation_id == variation_a]
        results_b = [r for r in self.test_results if r.variation_id == variation_b]
        
        if len(results_a) < min_samples or len(results_b) < min_samples:
            raise ValueError(f"Insufficient samples. Need at least {min_samples} for each variation.")
        
        # Extract metric values
        values_a = self._extract_metric_values(results_a, metric)
        values_b = self._extract_metric_values(results_b, metric)
        
        # Calculate statistics
        mean_a = np.mean(values_a)
        mean_b = np.mean(values_b)
        
        # Calculate improvement percentage
        improvement = ((mean_b - mean_a) / mean_a) * 100
        
        # Simple t-test (normally would use scipy.stats.ttest_ind)
        # For now, using basic statistical comparison
        std_a = np.std(values_a)
        std_b = np.std(values_b)
        n_a = len(values_a)
        n_b = len(values_b)
        
        # Basic p-value estimation (simplified)
        pooled_std = np.sqrt(((n_a - 1) * std_a**2 + (n_b - 1) * std_b**2) / (n_a + n_b - 2))
        t_stat = (mean_b - mean_a) / (pooled_std * np.sqrt(1/n_a + 1/n_b))
        p_value = 2 * (1 - abs(t_stat) / (abs(t_stat) + 1))  # Rough approximation
        
        # Confidence interval (simplified)
        margin_of_error = 1.96 * pooled_std * np.sqrt(1/n_a + 1/n_b)
        ci_lower = (mean_b - mean_a) - margin_of_error
        ci_upper = (mean_b - mean_a) + margin_of_error
        
        # Determine significance
        is_significant = p_value < 0.05
        
        # Generate recommendation
        if is_significant:
            if improvement > 0:
                recommendation = f"Variation B shows significant improvement ({improvement:.1f}%). Recommend adopting variation B."
            else:
                recommendation = f"Variation A shows significant improvement ({-improvement:.1f}%). Recommend keeping variation A."
        else:
            recommendation = f"No significant difference detected (p={p_value:.3f}). More data needed or variations are equivalent."
        
        return ComparisonResult(
            variation_a=variation_a,
            variation_b=variation_b,
            metric=metric,
            a_mean=mean_a,
            b_mean=mean_b,
            improvement_percent=improvement,
            p_value=p_value,
            confidence_interval=(ci_lower, ci_upper),
            is_significant=is_significant,
            sample_size=min(n_a, n_b),
            recommendation=recommendation
        )
    
    def get_best_variation(
        self,
        query_type: QueryType,
        metric: OptimizationMetric,
        min_samples: int = 10
    ) -> Optional[str]:
        """
        Get the best performing variation for a query type and metric.
        
        Args:
            query_type: Type of query
            metric: Metric to optimize for
            min_samples: Minimum samples required
            
        Returns:
            Best variation ID or None if insufficient data
        """
        # Filter results by query type
        relevant_results = [r for r in self.test_results if r.query_type == query_type]
        
        # Group by variation
        variation_performance = defaultdict(list)
        for result in relevant_results:
            variation_performance[result.variation_id].append(result)
        
        # Calculate mean performance for each variation
        best_variation = None
        best_score = None
        
        for variation_id, results in variation_performance.items():
            if len(results) >= min_samples:
                values = self._extract_metric_values(results, metric)
                mean_score = np.mean(values)
                
                if best_score is None or mean_score > best_score:
                    best_score = mean_score
                    best_variation = variation_id
        
        return best_variation
    
    def generate_optimization_report(
        self,
        experiment_id: str,
        output_file: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Generate a comprehensive optimization report.
        
        Args:
            experiment_id: Experiment to analyze
            output_file: Optional file to save report
            
        Returns:
            Report dictionary
        """
        if experiment_id not in self.active_experiments:
            raise ValueError(f"Experiment {experiment_id} not found")
        
        variation_ids = self.active_experiments[experiment_id]
        experiment_results = [r for r in self.test_results if r.variation_id in variation_ids]
        
        if not experiment_results:
            raise ValueError(f"No results found for experiment {experiment_id}")
        
        # Analyze each metric
        metrics = [
            OptimizationMetric.RESPONSE_TIME,
            OptimizationMetric.CONFIDENCE_SCORE,
            OptimizationMetric.CITATION_COUNT,
            OptimizationMetric.ANSWER_LENGTH
        ]
        
        report = {
            "experiment_id": experiment_id,
            "variations_tested": len(variation_ids),
            "total_tests": len(experiment_results),
            "analysis_date": time.time(),
            "metric_analysis": {},
            "recommendations": []
        }
        
        # Analyze each metric across variations
        for metric in metrics:
            metric_data = {}
            for variation_id in variation_ids:
                var_results = [r for r in experiment_results if r.variation_id == variation_id]
                if var_results:
                    values = self._extract_metric_values(var_results, metric)
                    metric_data[variation_id] = {
                        "mean": np.mean(values),
                        "std": np.std(values),
                        "count": len(values)
                    }
            
            report["metric_analysis"][metric.value] = metric_data
        
        # Generate recommendations
        for metric in metrics:
            best_variation = self.get_best_variation(
                query_type=QueryType.GENERAL,  # Could be made more specific
                metric=metric,
                min_samples=5
            )
            if best_variation:
                report["recommendations"].append({
                    "metric": metric.value,
                    "best_variation": best_variation,
                    "variation_name": self.variations[best_variation].name
                })
        
        # Save report if requested
        if output_file:
            with open(output_file, 'w') as f:
                json.dump(report, f, indent=2)
        
        return report
    
    def _extract_metric_values(self, results: List[TestResult], metric: OptimizationMetric) -> List[float]:
        """Extract metric values from test results."""
        values = []
        for result in results:
            if metric == OptimizationMetric.RESPONSE_TIME:
                values.append(result.response_time)
            elif metric == OptimizationMetric.CONFIDENCE_SCORE:
                values.append(result.confidence_score)
            elif metric == OptimizationMetric.CITATION_COUNT:
                values.append(float(result.citation_count))
            elif metric == OptimizationMetric.ANSWER_LENGTH:
                values.append(float(result.answer_length))
            elif metric == OptimizationMetric.TECHNICAL_ACCURACY and result.technical_accuracy is not None:
                values.append(result.technical_accuracy)
            elif metric == OptimizationMetric.USER_SATISFACTION and result.user_satisfaction is not None:
                values.append(result.user_satisfaction)
        
        return values
    
    def _load_experiments(self) -> None:
        """Load existing experiments from disk."""
        if not self.experiment_dir.exists():
            return
        
        for file_path in self.experiment_dir.glob("*.json"):
            if file_path.name.startswith("exp_"):
                with open(file_path, 'r') as f:
                    config = json.load(f)
                    self.active_experiments[config["experiment_id"]] = config["variation_ids"]
        
        # Load variations and results
        for file_path in self.experiment_dir.glob("variation_*.json"):
            with open(file_path, 'r') as f:
                var_data = json.load(f)
                variation = PromptVariation(**var_data)
                self.variations[variation.variation_id] = variation
        
        for file_path in self.experiment_dir.glob("result_*.json"):
            with open(file_path, 'r') as f:
                result_data = json.load(f)
                result = TestResult(**result_data)
                self.test_results.append(result)
    
    def _save_variation(self, variation: PromptVariation) -> None:
        """Save variation to disk."""
        file_path = self.experiment_dir / f"variation_{variation.variation_id}.json"
        var_dict = asdict(variation)
        
        # Convert template to dict
        var_dict["template"] = asdict(variation.template)
        var_dict["query_type"] = variation.query_type.value
        
        with open(file_path, 'w') as f:
            json.dump(var_dict, f, indent=2)
    
    def _save_test_result(self, result: TestResult) -> None:
        """Save test result to disk."""
        file_path = self.experiment_dir / f"result_{int(result.timestamp)}.json"
        result_dict = asdict(result)
        result_dict["query_type"] = result.query_type.value
        
        with open(file_path, 'w') as f:
            json.dump(result_dict, f, indent=2)


# Example usage
if __name__ == "__main__":
    # Initialize optimizer
    optimizer = PromptOptimizer()
    
    # Create temperature variations for implementation queries
    temp_variations = optimizer.create_temperature_variations(
        base_query_type=QueryType.IMPLEMENTATION,
        temperatures=[0.3, 0.7]
    )
    
    # Create length variations for definition queries
    length_variations = optimizer.create_length_variations(
        base_query_type=QueryType.DEFINITION,
        length_styles=["concise", "detailed"]
    )
    
    # Setup experiment
    test_queries = [
        "How do I implement a timer interrupt in RISC-V?",
        "What is the difference between machine mode and user mode?",
        "Configure GPIO pins for input/output operations"
    ]
    
    experiment_id = optimizer.setup_experiment(
        experiment_name="temperature_vs_length",
        variation_ids=temp_variations + length_variations,
        test_queries=test_queries
    )
    
    print(f"Created experiment: {experiment_id}")
    print(f"Variations: {len(temp_variations + length_variations)}")
    print(f"Test queries: {len(test_queries)}")