Upload 13 files

agent_tuner.py

@@ -0,0 +1,459 @@
+"""
+Agent Tuning Module for Agent Tuning Optimization Framework
+
+This module provides functionality for efficiently tuning large language models
+into specialized agents using a combination of positive examples, negative examples,
+and synthetically generated interaction trajectories.
+"""
+
+import os
+import torch
+import numpy as np
+from typing import List, Dict, Any, Union, Optional, Tuple
+from tqdm import tqdm
+from transformers import (
+    Trainer, TrainingArguments, 
+    DataCollatorForLanguageModeling,
+    AutoModelForCausalLM, AutoTokenizer
+)
+from datasets import Dataset
+
+from data.trajectory_data import Trajectory, TrajectoryDataset
+from models.llm_interface import LLMInterface
+
+class AgentTuner:
+    """Base class for agent tuning methods."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the agent tuner.
+        
+        Args:
+            name: Name of the tuning method
+        """
+        self.name = name
+    
+    def tune(
+        self, 
+        model_name: str,
+        trajectories: List[Trajectory],
+        **kwargs
+    ) -> Tuple[Any, Dict[str, Any]]:
+        """
+        Tune a model into a specialized agent.
+        
+        Args:
+            model_name: Name of the base model
+            trajectories: List of training trajectories
+            **kwargs: Additional tuning parameters
+            
+        Returns:
+            Tuple of (tuned_model, training_metrics)
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def save_model(self, model: Any, path: str) -> None:
+        """
+        Save the tuned model.
+        
+        Args:
+            model: Tuned model
+            path: Path to save the model
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def load_model(self, path: str) -> Any:
+        """
+        Load a tuned model.
+        
+        Args:
+            path: Path to the model
+            
+        Returns:
+            Loaded model
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class SupervisedFineTuner(AgentTuner):
+    """Tune agents using supervised fine-tuning."""
+    
+    def __init__(self):
+        """Initialize the supervised fine-tuner."""
+        super().__init__("supervised_fine_tuning")
+    
+    def tune(
+        self, 
+        model_name: str,
+        trajectories: List[Trajectory],
+        output_dir: str = "./tuned_model",
+        num_train_epochs: int = 3,
+        learning_rate: float = 5e-5,
+        batch_size: int = 4,
+        gradient_accumulation_steps: int = 4,
+        max_seq_length: int = 512,
+        format_type: str = "interleaved",
+        positive_weight: float = 0.8,
+        device: str = "cuda" if torch.cuda.is_available() else "cpu",
+        **kwargs
+    ) -> Tuple[Any, Dict[str, Any]]:
+        """
+        Tune a model using supervised fine-tuning.
+        
+        Args:
+            model_name: Name of the base model
+            trajectories: List of training trajectories
+            output_dir: Directory to save the model
+            num_train_epochs: Number of training epochs
+            learning_rate: Learning rate
+            batch_size: Batch size
+            gradient_accumulation_steps: Gradient accumulation steps
+            max_seq_length: Maximum sequence length
+            format_type: Format type for trajectories
+            positive_weight: Weight for positive examples
+            device: Device to use for training
+            **kwargs: Additional tuning parameters
+            
+        Returns:
+            Tuple of (tuned_model, training_metrics)
+        """
+        print(f"Starting supervised fine-tuning of {model_name}")
+        
+        # Create output directory
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Load model and tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        
+        # Ensure the tokenizer has a pad token
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Prepare training data
+        print("Preparing training data...")
+        
+        # Separate positive and negative trajectories
+        positive_trajectories = [t for t in trajectories if t.is_positive]
+        negative_trajectories = [t for t in trajectories if not t.is_positive]
+        
+        print(f"Found {len(positive_trajectories)} positive and {len(negative_trajectories)} negative trajectories")
+        
+        # Calculate sample counts based on positive weight
+        total_samples = len(trajectories)
+        target_positive = int(total_samples * positive_weight)
+        target_negative = total_samples - target_positive
+        
+        # Sample trajectories to achieve desired ratio
+        if len(positive_trajectories) > target_positive:
+            positive_trajectories = np.random.choice(positive_trajectories, target_positive, replace=False).tolist()
+        
+        if len(negative_trajectories) > target_negative:
+            negative_trajectories = np.random.choice(negative_trajectories, target_negative, replace=False).tolist()
+        
+        # Combine trajectories
+        sampled_trajectories = positive_trajectories + negative_trajectories
+        np.random.shuffle(sampled_trajectories)
+        
+        print(f"Using {len(positive_trajectories)} positive and {len(negative_trajectories)} negative trajectories for training")
+        
+        # Format trajectories for training
+        training_texts = []
+        
+        for trajectory in tqdm(sampled_trajectories, desc="Formatting trajectories"):
+            formatted = trajectory.to_training_format(format_type)
+            training_texts.append(formatted)
+        
+        # Tokenize training data
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=max_seq_length
+            )
+        
+        # Create dataset
+        dataset = Dataset.from_dict({"text": training_texts})
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["text"]
+        )
+        
+        # Set up training arguments
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=num_train_epochs,
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            learning_rate=learning_rate,
+            weight_decay=0.01,
+            save_strategy="epoch",
+            save_total_limit=2,
+            logging_dir=f"{output_dir}/logs",
+            logging_steps=10,
+            report_to="none"
+        )
+        
+        # Create data collator
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
+        
+        # Create trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset,
+            data_collator=data_collator
+        )
+        
+        # Train the model
+        print("Starting training...")
+        train_result = trainer.train()
+        
+        # Save the model
+        print(f"Saving model to {output_dir}")
+        trainer.save_model(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        
+        # Return the model and metrics
+        metrics = {
+            "train_loss": train_result.training_loss,
+            "train_runtime": train_result.metrics["train_runtime"],
+            "samples_per_second": train_result.metrics["train_samples_per_second"],
+            "num_train_samples": len(tokenized_dataset)
+        }
+        
+        return model, metrics
+    
+    def save_model(self, model: Any, path: str) -> None:
+        """
+        Save the tuned model.
+        
+        Args:
+            model: Tuned model
+            path: Path to save the model
+        """
+        model.save_pretrained(path)
+    
+    def load_model(self, path: str) -> Any:
+        """
+        Load a tuned model.
+        
+        Args:
+            path: Path to the model
+            
+        Returns:
+            Loaded model
+        """
+        return AutoModelForCausalLM.from_pretrained(path)
+
+
+class ParameterEfficientFineTuner(AgentTuner):
+    """Tune agents using parameter-efficient fine-tuning methods."""
+    
+    def __init__(self):
+        """Initialize the parameter-efficient fine-tuner."""
+        super().__init__("parameter_efficient_fine_tuning")
+    
+    def tune(
+        self, 
+        model_name: str,
+        trajectories: List[Trajectory],
+        output_dir: str = "./tuned_model",
+        method: str = "lora",  # 'lora', 'prefix', 'prompt_tuning'
+        num_train_epochs: int = 3,
+        learning_rate: float = 1e-4,
+        batch_size: int = 4,
+        gradient_accumulation_steps: int = 4,
+        max_seq_length: int = 512,
+        format_type: str = "interleaved",
+        positive_weight: float = 0.8,
+        device: str = "cuda" if torch.cuda.is_available() else "cpu",
+        **kwargs
+    ) -> Tuple[Any, Dict[str, Any]]:
+        """
+        Tune a model using parameter-efficient methods.
+        
+        Args:
+            model_name: Name of the base model
+            trajectories: List of training trajectories
+            output_dir: Directory to save the model
+            method: PEFT method to use
+            num_train_epochs: Number of training epochs
+            learning_rate: Learning rate
+            batch_size: Batch size
+            gradient_accumulation_steps: Gradient accumulation steps
+            max_seq_length: Maximum sequence length
+            format_type: Format type for trajectories
+            positive_weight: Weight for positive examples
+            device: Device to use for training
+            **kwargs: Additional tuning parameters
+            
+        Returns:
+            Tuple of (tuned_model, training_metrics)
+        """
+        try:
+            from peft import (
+                get_peft_model, LoraConfig, PrefixTuningConfig, 
+                PromptTuningConfig, TaskType, PeftModel
+            )
+        except ImportError:
+            raise ImportError("PEFT library is required for parameter-efficient fine-tuning. Install it with 'pip install peft'.")
+        
+        print(f"Starting parameter-efficient fine-tuning of {model_name} using {method}")
+        
+        # Create output directory
+        os.makedirs(output_dir, exist_ok=True)
+        
+        # Load model and tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        
+        # Ensure the tokenizer has a pad token
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Configure PEFT method
+        if method == "lora":
+            peft_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                r=16,
+                lora_alpha=32,
+                lora_dropout=0.1,
+                target_modules=["q_proj", "v_proj"]
+            )
+        elif method == "prefix":
+            peft_config = PrefixTuningConfig(
+                task_type=TaskType.CAUSAL_LM,
+                num_virtual_tokens=20,
+                prefix_projection=True
+            )
+        elif method == "prompt_tuning":
+            peft_config = PromptTuningConfig(
+                task_type=TaskType.CAUSAL_LM,
+                num_virtual_tokens=20,
+                tokenizer_name_or_path=model_name
+            )
+        else:
+            raise ValueError(f"Unsupported PEFT method: {method}")
+        
+        # Create PEFT model
+        model = get_peft_model(model, peft_config)
+        model.print_trainable_parameters()
+        
+        # Prepare training data (same as SupervisedFineTuner)
+        print("Preparing training data...")
+        
+        # Separate positive and negative trajectories
+        positive_trajectories = [t for t in trajectories if t.is_positive]
+        negative_trajectories = [t for t in trajectories if not t.is_positive]
+        
+        print(f"Found {len(positive_trajectories)} positive and {len(negative_trajectories)} negative trajectories")
+        
+        # Calculate sample counts based on positive weight
+        total_samples = len(trajectories)
+        target_positive = int(total_samples * positive_weight)
+        target_negative = total_samples - target_positive
+        
+        # Sample trajectories to achieve desired ratio
+        if len(positive_trajectories) > target_positive:
+            positive_trajectories = np.random.choice(positive_trajectories, target_positive, replace=False).tolist()
+        
+        if len(negative_trajectories) > target_negative:
+            negative_trajectories = np.random.choice(negative_trajectories, target_negative, replace=False).tolist()
+        
+        # Combine trajectories
+        sampled_trajectories = positive_trajectories + negative_trajectories
+        np.random.shuffle(sampled_trajectories)
+        
+        print(f"Using {len(positive_trajectories)} positive and {len(negative_trajectories)} negative trajectories for training")
+        
+        # Format trajectories for training
+        training_texts = []
+        
+        for trajectory in tqdm(sampled_trajectories, desc="Formatting trajectories"):
+            formatted = trajectory.to_training_format(format_type)
+            training_texts.append(formatted)
+        
+        # Tokenize training data
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=max_seq_length
+            )
+        
+        # Create dataset
+        dataset = Dataset.from_dict({"text": training_texts})
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["text"]
+        )
+        
+        # Set up training arguments
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=num_train_epochs,
+            per_device_train_batch_size=batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            learning_rate=learning_rate,
+            weight_decay=0.01,
+            save_strategy="epoch",
+            save_total_limit=2,
+            logging_dir=f"{output_dir}/logs",
+            logging_steps=10,
+            report_to="none"
+        )
+        
+        # Create data collator
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
+        
+        # Create trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset,
+            data_collator=data_collator
+        )
+        
+        # Train the model
+        print("Starting training...")
+        train_result = trainer.train()
+        
+        # Save the model
+        print(f"Saving model to {output_dir}")
+        trainer.save_model(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        
+        # Return the model and metrics
+        metrics = {
+            "train_loss": train_result.training_loss,
+            "train_runtime": train_result.metrics["train_runtime"],
+            "samples_per_second": train_result.metrics["train_samples_per_second"],
+            "num_train_samples": len(tokenized_dataset),
+            "peft_method": method
+        }
+        
+        return model, metrics
+    
+    def save_model(self, model: Any, path: str) -> None:
+        """
+        Save the tuned model.
+        
+        Args:
+            model: Tuned model
+            path: Path to save the model
+        """
+        model.save_pretrained(path)
+    
+  
+(Content truncated due to size limit. Use line ranges to read in chunks)

calibrators.py

@@ -0,0 +1,299 @@
+"""
+Domain-Specific Calibration Module for LLMs
+
+This module implements calibration techniques for improving uncertainty estimates
+across different domains, focusing on temperature scaling and domain adaptation.
+"""
+
+import numpy as np
+import torch
+from typing import List, Dict, Any, Union, Optional, Tuple
+from scipy.optimize import minimize_scalar
+
+class Calibrator:
+    """Base class for calibration methods."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the calibrator.
+        
+        Args:
+            name: Name of the calibration method
+        """
+        self.name = name
+        self.is_fitted = False
+    
+    def fit(self, confidences: List[float], accuracies: List[bool]) -> None:
+        """
+        Fit the calibrator to the provided data.
+        
+        Args:
+            confidences: List of confidence scores
+            accuracies: List of boolean accuracy indicators
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def calibrate(self, confidences: List[float]) -> List[float]:
+        """
+        Calibrate the provided confidence scores.
+        
+        Args:
+            confidences: List of confidence scores
+            
+        Returns:
+            Calibrated confidence scores
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class TemperatureScaling(Calibrator):
+    """Calibration using temperature scaling."""
+    
+    def __init__(self):
+        """Initialize the temperature scaling calibrator."""
+        super().__init__("temperature_scaling")
+        self.temperature = 1.0
+    
+    def _nll_loss(self, temperature: float, confidences: np.ndarray, accuracies: np.ndarray) -> float:
+        """
+        Calculate negative log likelihood loss for temperature scaling.
+        
+        Args:
+            temperature: Temperature parameter
+            confidences: Array of confidence scores
+            accuracies: Array of boolean accuracy indicators
+            
+        Returns:
+            Negative log likelihood loss
+        """
+        # Apply temperature scaling
+        scaled_confidences = np.clip(confidences / temperature, 1e-10, 1.0 - 1e-10)
+        
+        # Calculate binary cross-entropy loss
+        loss = -np.mean(
+            accuracies * np.log(scaled_confidences) + 
+            (1 - accuracies) * np.log(1 - scaled_confidences)
+        )
+        
+        return loss
+    
+    def fit(self, confidences: List[float], accuracies: List[bool]) -> None:
+        """
+        Fit the temperature parameter to the provided data.
+        
+        Args:
+            confidences: List of confidence scores
+            accuracies: List of boolean accuracy indicators
+        """
+        if not confidences or len(confidences) != len(accuracies):
+            raise ValueError("Confidences and accuracies must have the same non-zero length")
+        
+        # Convert to numpy arrays
+        conf_array = np.array(confidences)
+        acc_array = np.array(accuracies, dtype=float)
+        
+        # Optimize temperature parameter
+        result = minimize_scalar(
+            lambda t: self._nll_loss(t, conf_array, acc_array),
+            bounds=(0.1, 10.0),
+            method='bounded'
+        )
+        
+        self.temperature = result.x
+        self.is_fitted = True
+        
+        print(f"Fitted temperature parameter: {self.temperature:.4f}")
+    
+    def calibrate(self, confidences: List[float]) -> List[float]:
+        """
+        Calibrate the provided confidence scores using temperature scaling.
+        
+        Args:
+            confidences: List of confidence scores
+            
+        Returns:
+            Calibrated confidence scores
+        """
+        if not self.is_fitted:
+            raise ValueError("Calibrator must be fitted before calibration")
+        
+        # Apply temperature scaling
+        calibrated = [min(max(conf / self.temperature, 1e-10), 1.0 - 1e-10) for conf in confidences]
+        
+        return calibrated
+
+
+class DomainAdaptiveCalibration(Calibrator):
+    """Calibration using domain-adaptive techniques."""
+    
+    def __init__(self, source_domain: str, target_domain: str):
+        """
+        Initialize the domain-adaptive calibrator.
+        
+        Args:
+            source_domain: Source domain name
+            target_domain: Target domain name
+        """
+        super().__init__("domain_adaptive_calibration")
+        self.source_domain = source_domain
+        self.target_domain = target_domain
+        self.source_temperature = 1.0
+        self.target_temperature = 1.0
+        self.domain_shift_factor = 1.0
+    
+    def fit(
+        self, 
+        source_confidences: List[float], 
+        source_accuracies: List[bool],
+        target_confidences: Optional[List[float]] = None,
+        target_accuracies: Optional[List[bool]] = None
+    ) -> None:
+        """
+        Fit the domain-adaptive calibrator to the provided data.
+        
+        Args:
+            source_confidences: List of confidence scores from source domain
+            source_accuracies: List of boolean accuracy indicators from source domain
+            target_confidences: List of confidence scores from target domain (if available)
+            target_accuracies: List of boolean accuracy indicators from target domain (if available)
+        """
+        # Fit source domain temperature
+        source_calibrator = TemperatureScaling()
+        source_calibrator.fit(source_confidences, source_accuracies)
+        self.source_temperature = source_calibrator.temperature
+        
+        # If target domain data is available, fit target temperature
+        if target_confidences and target_accuracies:
+            target_calibrator = TemperatureScaling()
+            target_calibrator.fit(target_confidences, target_accuracies)
+            self.target_temperature = target_calibrator.temperature
+            
+            # Calculate domain shift factor
+            self.domain_shift_factor = self.target_temperature / self.source_temperature
+        else:
+            # Default domain shift factor based on heuristics
+            # This is a simplified approach; in a real system, this would be more sophisticated
+            self.domain_shift_factor = 1.2  # Assuming target domain is slightly more uncertain
+            self.target_temperature = self.source_temperature * self.domain_shift_factor
+        
+        self.is_fitted = True
+        
+        print(f"Fitted source temperature: {self.source_temperature:.4f}")
+        print(f"Fitted target temperature: {self.target_temperature:.4f}")
+        print(f"Domain shift factor: {self.domain_shift_factor:.4f}")
+    
+    def calibrate(self, confidences: List[float], domain: str = None) -> List[float]:
+        """
+        Calibrate the provided confidence scores using domain-adaptive calibration.
+        
+        Args:
+            confidences: List of confidence scores
+            domain: Domain of the confidences ('source' or 'target', defaults to target)
+            
+        Returns:
+            Calibrated confidence scores
+        """
+        if not self.is_fitted:
+            raise ValueError("Calibrator must be fitted before calibration")
+        
+        # Determine which temperature to use
+        if domain == "source":
+            temperature = self.source_temperature
+        else:
+            temperature = self.target_temperature
+        
+        # Apply temperature scaling
+        calibrated = [min(max(conf / temperature, 1e-10), 1.0 - 1e-10) for conf in confidences]
+        
+        return calibrated
+
+
+class EnsembleCalibration(Calibrator):
+    """Calibration using an ensemble of calibration methods."""
+    
+    def __init__(self, calibrators: List[Calibrator], weights: Optional[List[float]] = None):
+        """
+        Initialize the ensemble calibrator.
+        
+        Args:
+            calibrators: List of calibrator instances
+            weights: List of weights for each calibrator (None for equal weights)
+        """
+        super().__init__("ensemble_calibration")
+        self.calibrators = calibrators
+        
+        # Initialize weights
+        if weights is None:
+            self.weights = [1.0 / len(calibrators)] * len(calibrators)
+        else:
+            if len(weights) != len(calibrators):
+                raise ValueError("Number of weights must match number of calibrators")
+            
+            # Normalize weights
+            total = sum(weights)
+            self.weights = [w / total for w in weights]
+    
+    def fit(self, confidences: List[float], accuracies: List[bool]) -> None:
+        """
+        Fit all calibrators in the ensemble.
+        
+        Args:
+            confidences: List of confidence scores
+            accuracies: List of boolean accuracy indicators
+        """
+        for calibrator in self.calibrators:
+            calibrator.fit(confidences, accuracies)
+        
+        self.is_fitted = True
+    
+    def calibrate(self, confidences: List[float]) -> List[float]:
+        """
+        Calibrate the provided confidence scores using the ensemble.
+        
+        Args:
+            confidences: List of confidence scores
+            
+        Returns:
+            Calibrated confidence scores
+        """
+        if not self.is_fitted:
+            raise ValueError("Calibrator must be fitted before calibration")
+        
+        # Get calibrated confidences from each calibrator
+        all_calibrated = []
+        for calibrator in self.calibrators:
+            all_calibrated.append(calibrator.calibrate(confidences))
+        
+        # Combine calibrated confidences using weights
+        calibrated = []
+        for i in range(len(confidences)):
+            weighted_sum = sum(self.weights[j] * all_calibrated[j][i] for j in range(len(self.calibrators)))
+            calibrated.append(weighted_sum)
+        
+        return calibrated
+
+
+# Factory function to create calibrators
+def create_calibrator(method: str, **kwargs) -> Calibrator:
+    """
+    Create a calibrator based on the specified method.
+    
+    Args:
+        method: Name of the calibration method
+        **kwargs: Additional arguments for the calibrator
+        
+    Returns:
+        Calibrator instance
+    """
+    if method == "temperature_scaling":
+        return TemperatureScaling()
+    elif method == "domain_adaptive":
+        if "source_domain" not in kwargs or "target_domain" not in kwargs:
+            raise ValueError("Domain-adaptive calibration requires source_domain and target_domain")
+        return DomainAdaptiveCalibration(kwargs["source_domain"], kwargs["target_domain"])
+    elif method == "ensemble":
+        if "calibrators" not in kwargs:
+            raise ValueError("Ensemble calibration requires a list of calibrators")
+        return EnsembleCalibration(kwargs["calibrators"], kwargs.get("weights"))
+    else:
+        raise ValueError(f"Unsupported calibration method: {method}")

demo_app.py

@@ -0,0 +1,220 @@
+"""
+Simplified Gradio Demo for Agent Tuning Optimization Framework
+
+This script creates a simple Gradio web interface to demonstrate the framework's capabilities.
+"""
+
+import os
+import gradio as gr
+import numpy as np
+import random
+from datetime import datetime
+
+# Mock functions to simulate framework behavior without requiring full model loading
+def mock_generate_response(task, user_message):
+    """Simulate generating a response from a tuned agent."""
+    responses = [
+        f"I'll help you with your task to {task.lower()}. Based on your message '{user_message}', I recommend starting by breaking this down into smaller steps.",
+        f"I understand you need assistance with {task.lower()}. From your message, I can see that you're looking for guidance on '{user_message}'. Here's my approach to solving this.",
+        f"Thank you for providing details about {task.lower()}. Your message '{user_message}' gives me enough context to help you effectively. Let me outline a solution.",
+        f"I'm analyzing your request about {task.lower()}. Your message '{user_message}' indicates you need comprehensive assistance. Here's what I suggest as next steps."
+    ]
+    
+    # Simulate processing time
+    import time
+    time.sleep(1.5)
+    
+    return random.choice(responses) + f"\n\nResponse generated at {datetime.now().strftime('%H:%M:%S')}"
+
+def mock_generate_negative_sample(task, user_message, agent_message):
+    """Simulate generating a negative sample from a positive example."""
+    degradation_types = [
+        "Response truncation",
+        "Grammatical errors",
+        "Task misalignment",
+        "Constraint violation",
+        "Irrelevant tangent"
+    ]
+    
+    degradation = random.choice(degradation_types)
+    
+    if degradation == "Response truncation":
+        words = agent_message.split()
+        truncate_point = int(len(words) * random.uniform(0.3, 0.7))
+        return " ".join(words[:truncate_point]) + f"...\n\nNegative sample type: {degradation}"
+    
+    elif degradation == "Grammatical errors":
+        errors = [
+            lambda t: t.replace(".", ""),  # Remove periods
+            lambda t: t.replace("I ", "i "),  # Lowercase I
+            lambda t: t.replace(" the ", " teh "),  # Typo
+            lambda t: t.replace(" is ", " are "),  # Grammar error
+            lambda t: t.replace(" are ", " is ")  # Grammar error
+        ]
+        
+        result = agent_message
+        for _ in range(random.randint(2, 4)):
+            error_func = random.choice(errors)
+            result = error_func(result)
+        
+        return result + f"\n\nNegative sample type: {degradation}"
+    
+    elif degradation == "Task misalignment":
+        misalignments = [
+            f"I understand you're asking about something completely different. Let me tell you about weather patterns instead.",
+            f"I don't think that's what you really want to know. Let me explain something else that might interest you.",
+            f"Your question seems to be about {task}, but I'd rather discuss the history of computing.",
+            f"Instead of addressing your specific request about {task}, let me give you general information that's only tangentially related."
+        ]
+        
+        return random.choice(misalignments) + f"\n\nNegative sample type: {degradation}"
+    
+    elif degradation == "Constraint violation":
+        violations = [
+            f"I specifically recommend the XYZ Pro 2000 for $499.99, the UltraBook 15 for $1,299, and the PowerTech 5000 for $799. These are the absolute best options available.",
+            f"The system utilizes a polymorphic encapsulation paradigm with recursive lambda functions and stochastic gradient descent with backpropagation through a multi-layer perceptron.",
+            f"What specific features are you looking for? Do you have any brand preferences? What's your budget range? When do you need this by? Have you considered alternative options?",
+            f"Since you're a tech-savvy individual who values cutting-edge features, you'll definitely want the latest model with all the advanced capabilities."
+        ]
+        
+        return random.choice(violations) + f"\n\nNegative sample type: {degradation}"
+    
+    else:  # Irrelevant tangent
+        tangents = [
+            f"Did you know that artificial intelligence has been a concept since the 1950s? The field has evolved significantly since then, with major breakthroughs in neural networks and deep learning.",
+            f"I've been thinking about the philosophical implications of consciousness in AI systems. The question of whether an AI can truly understand or merely simulate understanding is fascinating.",
+            f"The weather has been quite interesting lately, with unusual patterns emerging globally. Climate scientists attribute this to a combination of factors including ocean temperature changes.",
+            f"I recently processed some fascinating data about renewable energy technologies. Solar efficiency has improved dramatically in the past decade, while costs have decreased by over 80%."
+        ]
+        
+        return random.choice(tangents) + f"\n\nNegative sample type: {degradation}"
+
+def mock_generate_synthetic_trajectory(task):
+    """Simulate generating a synthetic trajectory for a given task."""
+    # Determine task category
+    categories = ["travel", "shopping", "technology", "education", "finance", "health", "career", "home"]
+    category = random.choice(categories)
+    
+    # Generate interactions (2-4 turns)
+    num_turns = random.randint(2, 4)
+    interactions = []
+    
+    for j in range(num_turns):
+        if j == 0:
+            user_msg = f"I need help with this task: {task}"
+            agent_msg = f"I'd be happy to help you {task.lower()}. Could you provide more details about your preferences?"
+        elif j == num_turns - 1:
+            user_msg = "That sounds good. Please proceed with the final steps."
+            agent_msg = f"I've completed the task to {task.lower()}. Here's a summary of what I did..."
+        else:
+            user_msg = f"I prefer options that are {['affordable', 'convenient', 'high-quality'][j % 3]}."
+            agent_msg = f"Based on your preference for {['affordable', 'convenient', 'high-quality'][j % 3]} options, I recommend..."
+        
+        interactions.append({
+            'user': user_msg,
+            'agent': agent_msg
+        })
+    
+    # Format trajectory
+    result = f"Synthetic Trajectory for Task: {task}\nCategory: {category}\n\n"
+    
+    for i, interaction in enumerate(interactions):
+        result += f"Turn {i+1}:\nUser: {interaction['user']}\nAgent: {interaction['agent']}\n\n"
+    
+    result += f"Generation method: Template-based\nQuality score: {random.uniform(0.7, 0.9):.2f}"
+    
+    return result
+
+# Create Gradio interface
+with gr.Blocks(title="Agent Tuning Framework Demo") as demo:
+    gr.Markdown("# Agent Tuning Optimization Framework Demo")
+    gr.Markdown("### A framework for efficiently tuning LLMs into specialized agents using negative and synthetic samples")
+    
+    with gr.Tab("Generate Response"):
+        with gr.Row():
+            with gr.Column():
+                task_input = gr.Textbox(label="Task Description", placeholder="e.g., Book a flight from New York to London")
+                user_input = gr.Textbox(label="User Message", placeholder="e.g., I need to travel next week for business")
+                generate_btn = gr.Button("Generate Response", variant="primary")
+            with gr.Column():
+                response_output = gr.Textbox(label="Agent Response", lines=8)
+        
+        generate_btn.click(
+            mock_generate_response,
+            inputs=[task_input, user_input],
+            outputs=response_output
+        )
+        
+        gr.Examples(
+            [
+                ["Book a flight from New York to London", "I need to travel next week for business"],
+                ["Find a vegetarian restaurant", "I'm looking for dinner options tonight"],
+                ["Help me debug a Python script", "I'm getting an IndexError in my code"]
+            ],
+            inputs=[task_input, user_input]
+        )
+    
+    with gr.Tab("Generate Negative Sample"):
+        with gr.Row():
+            with gr.Column():
+                neg_task_input = gr.Textbox(label="Task Description", placeholder="e.g., Book a flight from New York to London")
+                neg_user_input = gr.Textbox(label="User Message", placeholder="e.g., I need to travel next week for business")
+                neg_agent_input = gr.Textbox(label="Agent Message (Positive Example)", placeholder="e.g., I'd be happy to help you book a flight...", lines=5)
+                neg_generate_btn = gr.Button("Generate Negative Sample", variant="primary")
+            with gr.Column():
+                neg_output = gr.Textbox(label="Negative Sample", lines=8)
+        
+        neg_generate_btn.click(
+            mock_generate_negative_sample,
+            inputs=[neg_task_input, neg_user_input, neg_agent_input],
+            outputs=neg_output
+        )
+        
+        gr.Examples(
+            [
+                ["Book a flight from New York to London", "I need to travel next week for business", "I'd be happy to help you book a flight from New York to London. Could you provide more details about your preferred travel dates, budget, and any airline preferences you might have?"],
+                ["Recommend a laptop for programming", "I need a new laptop for software development", "I can definitely help you find a suitable laptop for programming. Based on software development needs, I'd recommend looking for a laptop with at least 16GB RAM, a multi-core processor, and an SSD for storage. Would you like specific brand recommendations or have a particular budget in mind?"]
+            ],
+            inputs=[neg_task_input, neg_user_input, neg_agent_input]
+        )
+    
+    with gr.Tab("Generate Synthetic Trajectory"):
+        with gr.Row():
+            with gr.Column():
+                synth_task_input = gr.Textbox(label="Task Description", placeholder="e.g., Plan a weekend trip to Chicago")
+                synth_generate_btn = gr.Button("Generate Synthetic Trajectory", variant="primary")
+            with gr.Column():
+                synth_output = gr.Textbox(label="Synthetic Trajectory", lines=15)
+        
+        synth_generate_btn.click(
+            mock_generate_synthetic_trajectory,
+            inputs=[synth_task_input],
+            outputs=synth_output
+        )
+        
+        gr.Examples(
+            [
+                ["Plan a weekend trip to Chicago"],
+                ["Recommend healthy meal prep options for the week"],
+                ["Help me create a study schedule for final exams"]
+            ],
+            inputs=[synth_task_input]
+        )
+    
+    gr.Markdown("""
+    ## About This Framework
+    
+    The Agent Tuning Optimization Framework provides a comprehensive solution for efficiently tuning large language models into specialized agents through the strategic incorporation of negative samples and synthetic trajectories.
+    
+    ### Key Features:
+    
+    1. **Negative Sample Generation**: Creates examples of undesired agent behaviors to teach models what not to do
+    2. **Synthetic Trajectory Generation**: Automatically generates diverse interaction trajectories
+    3. **Mixed-Sample Tuning**: Combines positive examples, negative samples, and synthetic trajectories
+    4. **Parameter-Efficient Fine-Tuning**: Implements methods like LoRA for computational efficiency
+    
+    This demo provides a simplified simulation of the framework's capabilities. For full functionality, deploy the complete framework following the provided documentation.
+    """)
+
+# Launch the interface
+demo.launch(share=True)

deploy.py

@@ -0,0 +1,354 @@
+"""
+Deployment Script for Agent Tuning Optimization Framework
+
+This script prepares the framework for deployment to production environments
+and Hugging Face Spaces.
+"""
+
+import os
+import shutil
+import argparse
+import subprocess
+import json
+from pathlib import Path
+
+def prepare_for_deployment(source_dir, output_dir, config_path=None):
+    """
+    Prepare the framework for deployment.
+    
+    Args:
+        source_dir: Source directory containing the framework
+        output_dir: Output directory for deployment package
+        config_path: Path to configuration file (optional)
+    """
+    print(f"Preparing deployment package from {source_dir} to {output_dir}")
+    
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Copy core modules
+    core_modules = [
+        "models",
+        "data",
+        "training",
+        "evaluation",
+        "main.py",
+        "README.md"
+    ]
+    
+    for module in core_modules:
+        source_path = os.path.join(source_dir, module)
+        target_path = os.path.join(output_dir, module)
+        
+        if os.path.isdir(source_path):
+            if os.path.exists(target_path):
+                shutil.rmtree(target_path)
+            shutil.copytree(source_path, target_path)
+        else:
+            shutil.copy2(source_path, target_path)
+    
+    # Copy configuration file if provided
+    if config_path:
+        shutil.copy2(config_path, os.path.join(output_dir, "config.json"))
+    else:
+        # Use example config
+        example_config_path = os.path.join(source_dir, "example_config.json")
+        if os.path.exists(example_config_path):
+            shutil.copy2(example_config_path, os.path.join(output_dir, "config.json"))
+    
+    # Create requirements.txt
+    requirements = [
+        "torch>=1.10.0",
+        "transformers>=4.20.0",
+        "datasets>=2.0.0",
+        "numpy>=1.20.0",
+        "pandas>=1.3.0",
+        "matplotlib>=3.4.0",
+        "tqdm>=4.60.0",
+        "scikit-learn>=1.0.0",
+        "peft>=0.2.0"
+    ]
+    
+    with open(os.path.join(output_dir, "requirements.txt"), "w") as f:
+        f.write("\n".join(requirements))
+    
+    # Create setup.py
+    setup_py = """
+from setuptools import setup, find_packages
+
+setup(
+    name="agent_tuning_framework",
+    version="0.1.0",
+    packages=find_packages(),
+    install_requires=[
+        "torch>=1.10.0",
+        "transformers>=4.20.0",
+        "datasets>=2.0.0",
+        "numpy>=1.20.0",
+        "pandas>=1.3.0",
+        "matplotlib>=3.4.0",
+        "tqdm>=4.60.0",
+        "scikit-learn>=1.0.0",
+        "peft>=0.2.0"
+    ],
+    author="MBZUAI Technical Interview Preparation",
+    author_email="example@example.com",
+    description="Agent Tuning Optimization Framework with Negative and Synthetic Samples",
+    keywords="nlp, machine learning, agent tuning, language models",
+    url="https://github.com/username/agent_tuning_framework",
+)
+"""
+    
+    with open(os.path.join(output_dir, "setup.py"), "w") as f:
+        f.write(setup_py)
+    
+    # Create app.py for web interface
+    app_py = """
+import os
+import json
+import gradio as gr
+import torch
+from models.llm_interface import LLMInterface
+from data.trajectory_data import TrajectoryDataset, Trajectory
+from training.negative_samples import create_negative_sample_generator
+from training.synthetic_trajectories import create_synthetic_trajectory_generator
+
+# Initialize model
+def load_model(model_path):
+    if os.path.exists(model_path):
+        return LLMInterface(
+            model_name=model_path,
+            model_type="causal",
+            device="cuda" if torch.cuda.is_available() else "cpu"
+        )
+    else:
+        return LLMInterface(
+            model_name="gpt2",
+            model_type="causal",
+            device="cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+# Initialize components
+model = load_model("./tuned_model")
+negative_generator = create_negative_sample_generator("response_degradation")
+synthetic_generator = create_synthetic_trajectory_generator("template")
+
+# Define interface functions
+def generate_response(task, user_message):
+    prompt = f"Task: {task}\\n\\nUser: {user_message}\\nAgent:"
+    response = model.generate(prompt)
+    return response["response"]
+
+def generate_negative_sample(task, user_message, agent_message):
+    trajectory = Trajectory(
+        task_description=task,
+        interactions=[{"user": user_message, "agent": agent_message}]
+    )
+    negative_trajectory = negative_generator.generate(trajectory)
+    return negative_trajectory.interactions[0]["agent"]
+
+def generate_synthetic_trajectory(task):
+    trajectory = synthetic_generator.generate(task)
+    result = ""
+    for i, interaction in enumerate(trajectory.interactions):
+        result += f"Turn {i+1}:\\nUser: {interaction['user']}\\nAgent: {interaction['agent']}\\n\\n"
+    return result
+
+# Create Gradio interface
+with gr.Blocks(title="Agent Tuning Framework Demo") as demo:
+    gr.Markdown("# Agent Tuning Optimization Framework Demo")
+    
+    with gr.Tab("Generate Response"):
+        with gr.Row():
+            with gr.Column():
+                task_input = gr.Textbox(label="Task Description")
+                user_input = gr.Textbox(label="User Message")
+                generate_btn = gr.Button("Generate Response")
+            with gr.Column():
+                response_output = gr.Textbox(label="Agent Response")
+        
+        generate_btn.click(
+            generate_response,
+            inputs=[task_input, user_input],
+            outputs=response_output
+        )
+    
+    with gr.Tab("Generate Negative Sample"):
+        with gr.Row():
+            with gr.Column():
+                neg_task_input = gr.Textbox(label="Task Description")
+                neg_user_input = gr.Textbox(label="User Message")
+                neg_agent_input = gr.Textbox(label="Agent Message (Positive Example)")
+                neg_generate_btn = gr.Button("Generate Negative Sample")
+            with gr.Column():
+                neg_output = gr.Textbox(label="Negative Sample")
+        
+        neg_generate_btn.click(
+            generate_negative_sample,
+            inputs=[neg_task_input, neg_user_input, neg_agent_input],
+            outputs=neg_output
+        )
+    
+    with gr.Tab("Generate Synthetic Trajectory"):
+        with gr.Row():
+            with gr.Column():
+                synth_task_input = gr.Textbox(label="Task Description")
+                synth_generate_btn = gr.Button("Generate Synthetic Trajectory")
+            with gr.Column():
+                synth_output = gr.Textbox(label="Synthetic Trajectory")
+        
+        synth_generate_btn.click(
+            generate_synthetic_trajectory,
+            inputs=[synth_task_input],
+            outputs=synth_output
+        )
+
+if __name__ == "__main__":
+    demo.launch()
+"""
+    
+    with open(os.path.join(output_dir, "app.py"), "w") as f:
+        f.write(app_py)
+    
+    # Create Dockerfile
+    dockerfile = """
+FROM python:3.9-slim
+
+WORKDIR /app
+
+COPY . /app/
+
+RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir gradio>=3.0.0
+
+EXPOSE 7860
+
+CMD ["python", "app.py"]
+"""
+    
+    with open(os.path.join(output_dir, "Dockerfile"), "w") as f:
+        f.write(dockerfile)
+    
+    # Create README for deployment
+    deployment_readme = """
+# Agent Tuning Optimization Framework
+
+This package contains the Agent Tuning Optimization Framework with Negative and Synthetic Samples, a comprehensive solution for efficiently tuning large language models into specialized agents.
+
+## Installation
+
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+
+### Running Experiments
+
+```bash
+python main.py --config config.json --output ./experiment_results
+```
+
+### Web Interface
+
+```bash
+pip install gradio
+python app.py
+```
+
+## Deployment Options
+
+### Docker
+
+```bash
+docker build -t agent-tuning-framework .
+docker run -p 7860:7860 agent-tuning-framework
+```
+
+### Hugging Face Spaces
+
+This project can be deployed to Hugging Face Spaces by following these steps:
+
+1. Create a new Space on Hugging Face (https://huggingface.co/spaces)
+2. Select "Gradio" as the SDK
+3. Upload all files from this directory to the Space
+4. The Space will automatically build and deploy the application
+
+## Configuration
+
+See `config.json` for configuration options.
+
+## License
+
+MIT
+"""
+    
+    with open(os.path.join(output_dir, "README.md"), "w") as f:
+        f.write(deployment_readme)
+    
+    # Create Hugging Face Space files
+    os.makedirs(os.path.join(output_dir, "huggingface"), exist_ok=True)
+    
+    # Create requirements.txt for Hugging Face
+    hf_requirements = requirements + ["gradio>=3.0.0"]
+    
+    with open(os.path.join(output_dir, "huggingface", "requirements.txt"), "w") as f:
+        f.write("\n".join(hf_requirements))
+    
+    # Copy app.py
+    shutil.copy2(os.path.join(output_dir, "app.py"), os.path.join(output_dir, "huggingface", "app.py"))
+    
+    # Create README for Hugging Face
+    hf_readme = """
+---
+title: Agent Tuning Optimization Framework
+emoji: 🤖
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 3.36.1
+app_file: app.py
+pinned: false
+license: mit
+---
+
+# Agent Tuning Optimization Framework
+
+This Space demonstrates the Agent Tuning Optimization Framework with Negative and Synthetic Samples, a comprehensive solution for efficiently tuning large language models into specialized agents.
+
+## Features
+
+- Generate agent responses for given tasks and user messages
+- Create negative samples from positive examples
+- Generate synthetic interaction trajectories
+
+## Usage
+
+1. Select a tab for the desired functionality
+2. Enter the required information
+3. Click the button to generate results
+
+## Learn More
+
+For more information, visit the [GitHub repository](https://github.com/username/agent_tuning_framework).
+"""
+    
+    with open(os.path.join(output_dir, "huggingface", "README.md"), "w") as f:
+        f.write(hf_readme)
+    
+    print(f"Deployment package prepared in {output_dir}")
+    print(f"Hugging Face Space files prepared in {os.path.join(output_dir, 'huggingface')}")
+
+def main():
+    """Main function for preparing deployment package."""
+    parser = argparse.ArgumentParser(description="Prepare deployment package for Agent Tuning Framework")
+    parser.add_argument("--source", type=str, default=".", help="Source directory containing the framework")
+    parser.add_argument("--output", type=str, default="./deployment", help="Output directory for deployment package")
+    parser.add_argument("--config", type=str, help="Path to configuration file")
+    
+    args = parser.parse_args()
+    
+    prepare_for_deployment(args.source, args.output, args.config)
+
+if __name__ == "__main__":
+    main()

domain_datasets.py

@@ -0,0 +1,364 @@
+"""
+Domain Dataset Module for Cross-Domain Uncertainty Quantification
+
+This module provides functionality for loading and managing datasets from different domains
+for evaluating uncertainty quantification methods across domains.
+"""
+
+import os
+import json
+import pandas as pd
+import numpy as np
+from typing import List, Dict, Any, Union, Optional, Tuple
+from datasets import load_dataset
+
+class DomainDataset:
+    """Base class for domain-specific datasets."""
+    
+    def __init__(self, name: str, domain: str):
+        """
+        Initialize the domain dataset.
+        
+        Args:
+            name: Name of the dataset
+            domain: Domain category (e.g., 'medical', 'legal', 'general')
+        """
+        self.name = name
+        self.domain = domain
+        self.data = None
+    
+    def load(self) -> None:
+        """Load the dataset."""
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def get_samples(self, n: Optional[int] = None) -> List[Dict[str, Any]]:
+        """
+        Get samples from the dataset.
+        
+        Args:
+            n: Number of samples to return (None for all)
+            
+        Returns:
+            List of samples with prompts and expected outputs
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def get_prompt_template(self) -> str:
+        """
+        Get the prompt template for this domain.
+        
+        Returns:
+            Prompt template string
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class MedicalQADataset(DomainDataset):
+    """Dataset for medical question answering."""
+    
+    def __init__(self, data_path: Optional[str] = None):
+        """
+        Initialize the medical QA dataset.
+        
+        Args:
+            data_path: Path to the dataset file (None to use default)
+        """
+        super().__init__("medical_qa", "medical")
+        self.data_path = data_path
+    
+    def load(self) -> None:
+        """Load the medical QA dataset."""
+        if self.data_path and os.path.exists(self.data_path):
+            # Load from local file if available
+            if self.data_path.endswith('.csv'):
+                self.data = pd.read_csv(self.data_path)
+            elif self.data_path.endswith('.json'):
+                with open(self.data_path, 'r') as f:
+                    self.data = json.load(f)
+            else:
+                raise ValueError(f"Unsupported file format: {self.data_path}")
+        else:
+            # Use a sample of the MedMCQA dataset from Hugging Face
+            try:
+                dataset = load_dataset("medmcqa", split="train[:100]")
+                self.data = dataset.to_pandas()
+            except Exception as e:
+                # Fallback to synthetic data if dataset loading fails
+                print(f"Failed to load MedMCQA dataset: {e}")
+                self.data = self._create_synthetic_data()
+    
+    def _create_synthetic_data(self) -> pd.DataFrame:
+        """Create synthetic medical QA data for testing."""
+        questions = [
+            "What are the common symptoms of myocardial infarction?",
+            "How does insulin regulate blood glucose levels?",
+            "What is the mechanism of action for ACE inhibitors?",
+            "What are the diagnostic criteria for rheumatoid arthritis?",
+            "How does the SARS-CoV-2 virus enter human cells?",
+            "What are the main side effects of chemotherapy?",
+            "How does the blood-brain barrier function?",
+            "What is the pathophysiology of type 2 diabetes?",
+            "How do vaccines create immunity?",
+            "What are the stages of chronic kidney disease?"
+        ]
+        
+        # Create a dataframe with questions only (answers would be generated by LLMs)
+        return pd.DataFrame({
+            'question': questions,
+            'domain': ['medical'] * len(questions)
+        })
+    
+    def get_samples(self, n: Optional[int] = None) -> List[Dict[str, Any]]:
+        """
+        Get samples from the medical QA dataset.
+        
+        Args:
+            n: Number of samples to return (None for all)
+            
+        Returns:
+            List of samples with prompts
+        """
+        if self.data is None:
+            self.load()
+        
+        if 'question' in self.data.columns:
+            questions = self.data['question'].tolist()
+        elif 'question_text' in self.data.columns:
+            questions = self.data['question_text'].tolist()
+        else:
+            raise ValueError("Dataset does not contain question column")
+        
+        if n is not None:
+            questions = questions[:n]
+        
+        # Create samples with prompts
+        samples = []
+        for question in questions:
+            prompt = self.get_prompt_template().format(question=question)
+            samples.append({
+                'domain': 'medical',
+                'question': question,
+                'prompt': prompt
+            })
+        
+        return samples
+    
+    def get_prompt_template(self) -> str:
+        """
+        Get the prompt template for medical domain.
+        
+        Returns:
+            Prompt template string
+        """
+        return "You are a medical expert. Please answer the following medical question accurately and concisely:\n\n{question}"
+
+
+class LegalQADataset(DomainDataset):
+    """Dataset for legal question answering."""
+    
+    def __init__(self, data_path: Optional[str] = None):
+        """
+        Initialize the legal QA dataset.
+        
+        Args:
+            data_path: Path to the dataset file (None to use default)
+        """
+        super().__init__("legal_qa", "legal")
+        self.data_path = data_path
+    
+    def load(self) -> None:
+        """Load the legal QA dataset."""
+        if self.data_path and os.path.exists(self.data_path):
+            # Load from local file if available
+            if self.data_path.endswith('.csv'):
+                self.data = pd.read_csv(self.data_path)
+            elif self.data_path.endswith('.json'):
+                with open(self.data_path, 'r') as f:
+                    self.data = json.load(f)
+            else:
+                raise ValueError(f"Unsupported file format: {self.data_path}")
+        else:
+            # Use synthetic data for legal domain
+            self.data = self._create_synthetic_data()
+    
+    def _create_synthetic_data(self) -> pd.DataFrame:
+        """Create synthetic legal QA data for testing."""
+        questions = [
+            "What constitutes a breach of contract?",
+            "How is intellectual property protected under international law?",
+            "What are the elements of negligence in tort law?",
+            "How does the doctrine of stare decisis function in common law systems?",
+            "What rights are protected under the Fourth Amendment?",
+            "What is the difference between a patent and a copyright?",
+            "How does arbitration differ from litigation?",
+            "What constitutes insider trading under securities law?",
+            "What are the legal requirements for a valid will?",
+            "How does diplomatic immunity work under international law?"
+        ]
+        
+        # Create a dataframe with questions only
+        return pd.DataFrame({
+            'question': questions,
+            'domain': ['legal'] * len(questions)
+        })
+    
+    def get_samples(self, n: Optional[int] = None) -> List[Dict[str, Any]]:
+        """
+        Get samples from the legal QA dataset.
+        
+        Args:
+            n: Number of samples to return (None for all)
+            
+        Returns:
+            List of samples with prompts
+        """
+        if self.data is None:
+            self.load()
+        
+        questions = self.data['question'].tolist()
+        
+        if n is not None:
+            questions = questions[:n]
+        
+        # Create samples with prompts
+        samples = []
+        for question in questions:
+            prompt = self.get_prompt_template().format(question=question)
+            samples.append({
+                'domain': 'legal',
+                'question': question,
+                'prompt': prompt
+            })
+        
+        return samples
+    
+    def get_prompt_template(self) -> str:
+        """
+        Get the prompt template for legal domain.
+        
+        Returns:
+            Prompt template string
+        """
+        return "You are a legal expert. Please answer the following legal question accurately and concisely:\n\n{question}"
+
+
+class GeneralKnowledgeDataset(DomainDataset):
+    """Dataset for general knowledge question answering."""
+    
+    def __init__(self, data_path: Optional[str] = None):
+        """
+        Initialize the general knowledge dataset.
+        
+        Args:
+            data_path: Path to the dataset file (None to use default)
+        """
+        super().__init__("general_knowledge", "general")
+        self.data_path = data_path
+    
+    def load(self) -> None:
+        """Load the general knowledge dataset."""
+        if self.data_path and os.path.exists(self.data_path):
+            # Load from local file if available
+            if self.data_path.endswith('.csv'):
+                self.data = pd.read_csv(self.data_path)
+            elif self.data_path.endswith('.json'):
+                with open(self.data_path, 'r') as f:
+                    self.data = json.load(f)
+            else:
+                raise ValueError(f"Unsupported file format: {self.data_path}")
+        else:
+            # Use a sample of the TriviaQA dataset from Hugging Face
+            try:
+                dataset = load_dataset("trivia_qa", "unfiltered", split="train[:100]")
+                self.data = dataset.to_pandas()
+            except Exception as e:
+                # Fallback to synthetic data if dataset loading fails
+                print(f"Failed to load TriviaQA dataset: {e}")
+                self.data = self._create_synthetic_data()
+    
+    def _create_synthetic_data(self) -> pd.DataFrame:
+        """Create synthetic general knowledge data for testing."""
+        questions = [
+            "What is the capital of France?",
+            "Who wrote the novel '1984'?",
+            "What is the chemical symbol for gold?",
+            "Which planet is known as the Red Planet?",
+            "Who painted the Mona Lisa?",
+            "What is the largest ocean on Earth?",
+            "What year did World War II end?",
+            "What is the tallest mountain in the world?",
+            "Who was the first person to step on the moon?",
+            "What is the speed of light in a vacuum?"
+        ]
+        
+        # Create a dataframe with questions only
+        return pd.DataFrame({
+            'question': questions,
+            'domain': ['general'] * len(questions)
+        })
+    
+    def get_samples(self, n: Optional[int] = None) -> List[Dict[str, Any]]:
+        """
+        Get samples from the general knowledge dataset.
+        
+        Args:
+            n: Number of samples to return (None for all)
+            
+        Returns:
+            List of samples with prompts
+        """
+        if self.data is None:
+            self.load()
+        
+        if 'question' in self.data.columns:
+            questions = self.data['question'].tolist()
+        elif 'question_text' in self.data.columns:
+            questions = self.data['question_text'].tolist()
+        else:
+            raise ValueError("Dataset does not contain question column")
+        
+        if n is not None:
+            questions = questions[:n]
+        
+        # Create samples with prompts
+        samples = []
+        for question in questions:
+            prompt = self.get_prompt_template().format(question=question)
+            samples.append({
+                'domain': 'general',
+                'question': question,
+                'prompt': prompt
+            })
+        
+        return samples
+    
+    def get_prompt_template(self) -> str:
+        """
+        Get the prompt template for general knowledge domain.
+        
+        Returns:
+            Prompt template string
+        """
+        return "Please answer the following general knowledge question accurately and concisely:\n\n{question}"
+
+
+# Factory function to create domain datasets
+def create_domain_dataset(domain: str, data_path: Optional[str] = None) -> DomainDataset:
+    """
+    Create a domain dataset based on the specified domain.
+    
+    Args:
+        domain: Domain category ('medical', 'legal', 'general')
+        data_path: Path to the dataset file (None to use default)
+        
+    Returns:
+        Domain dataset instance
+    """
+    if domain == "medical":
+        return MedicalQADataset(data_path)
+    elif domain == "legal":
+        return LegalQADataset(data_path)
+    elif domain == "general":
+        return GeneralKnowledgeDataset(data_path)
+    else:
+        raise ValueError(f"Unsupported domain: {domain}")

evaluators.py

@@ -0,0 +1,448 @@
+"""
+Evaluation Framework for Cross-Domain Uncertainty Quantification
+
+This module provides functionality for evaluating uncertainty quantification methods
+across different domains, including metrics for uncertainty quality and cross-domain performance.
+"""
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from typing import List, Dict, Any, Union, Optional, Tuple
+from sklearn.metrics import roc_auc_score, precision_recall_curve, auc
+
+class UncertaintyEvaluator:
+    """Evaluator for uncertainty quantification methods."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the uncertainty evaluator.
+        
+        Args:
+            name: Name of the evaluation method
+        """
+        self.name = name
+    
+    def evaluate(
+        self, 
+        uncertainties: List[float], 
+        correctness: List[bool]
+    ) -> Dict[str, float]:
+        """
+        Evaluate uncertainty estimates against correctness.
+        
+        Args:
+            uncertainties: List of uncertainty scores (higher means more uncertain)
+            correctness: List of boolean correctness indicators
+            
+        Returns:
+            Dictionary of evaluation metrics
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class CalibrationEvaluator(UncertaintyEvaluator):
+    """Evaluator for calibration quality."""
+    
+    def __init__(self):
+        """Initialize the calibration evaluator."""
+        super().__init__("calibration_evaluator")
+    
+    def expected_calibration_error(
+        self, 
+        confidences: List[float], 
+        correctness: List[bool], 
+        num_bins: int = 10
+    ) -> float:
+        """
+        Calculate Expected Calibration Error (ECE).
+        
+        Args:
+            confidences: List of confidence scores
+            correctness: List of boolean correctness indicators
+            num_bins: Number of bins for binning confidences
+            
+        Returns:
+            Expected Calibration Error
+        """
+        if len(confidences) != len(correctness):
+            raise ValueError("Confidences and correctness must have the same length")
+        
+        if not confidences:
+            return 0.0
+        
+        # Create bins and calculate ECE
+        bin_indices = np.digitize(confidences, np.linspace(0, 1, num_bins))
+        ece = 0.0
+        
+        for bin_idx in range(1, num_bins + 1):
+            bin_mask = (bin_indices == bin_idx)
+            if np.any(bin_mask):
+                bin_confidences = np.array(confidences)[bin_mask]
+                bin_correctness = np.array(correctness)[bin_mask]
+                bin_confidence = np.mean(bin_confidences)
+                bin_accuracy = np.mean(bin_correctness)
+                bin_size = np.sum(bin_mask)
+                
+                # Weighted absolute difference between confidence and accuracy
+                ece += (bin_size / len(confidences)) * np.abs(bin_confidence - bin_accuracy)
+        
+        return float(ece)
+    
+    def maximum_calibration_error(
+        self, 
+        confidences: List[float], 
+        correctness: List[bool], 
+        num_bins: int = 10
+    ) -> float:
+        """
+        Calculate Maximum Calibration Error (MCE).
+        
+        Args:
+            confidences: List of confidence scores
+            correctness: List of boolean correctness indicators
+            num_bins: Number of bins for binning confidences
+            
+        Returns:
+            Maximum Calibration Error
+        """
+        if len(confidences) != len(correctness):
+            raise ValueError("Confidences and correctness must have the same length")
+        
+        if not confidences:
+            return 0.0
+        
+        # Create bins and calculate MCE
+        bin_indices = np.digitize(confidences, np.linspace(0, 1, num_bins))
+        max_ce = 0.0
+        
+        for bin_idx in range(1, num_bins + 1):
+            bin_mask = (bin_indices == bin_idx)
+            if np.any(bin_mask):
+                bin_confidences = np.array(confidences)[bin_mask]
+                bin_correctness = np.array(correctness)[bin_mask]
+                bin_confidence = np.mean(bin_confidences)
+                bin_accuracy = np.mean(bin_correctness)
+                
+                # Absolute difference between confidence and accuracy
+                ce = np.abs(bin_confidence - bin_accuracy)
+                max_ce = max(max_ce, ce)
+        
+        return float(max_ce)
+    
+    def evaluate(
+        self, 
+        confidences: List[float], 
+        correctness: List[bool]
+    ) -> Dict[str, float]:
+        """
+        Evaluate calibration quality.
+        
+        Args:
+            confidences: List of confidence scores
+            correctness: List of boolean correctness indicators
+            
+        Returns:
+            Dictionary of calibration metrics:
+                - ece: Expected Calibration Error
+                - mce: Maximum Calibration Error
+        """
+        return {
+            "ece": self.expected_calibration_error(confidences, correctness),
+            "mce": self.maximum_calibration_error(confidences, correctness)
+        }
+    
+    def plot_reliability_diagram(
+        self, 
+        confidences: List[float], 
+        correctness: List[bool], 
+        num_bins: int = 10,
+        title: str = "Reliability Diagram",
+        save_path: Optional[str] = None
+    ) -> None:
+        """
+        Plot a reliability diagram for calibration visualization.
+        
+        Args:
+            confidences: List of confidence scores
+            correctness: List of boolean correctness indicators
+            num_bins: Number of bins for binning confidences
+            title: Title for the plot
+            save_path: Path to save the plot (None to display)
+        """
+        if len(confidences) != len(correctness):
+            raise ValueError("Confidences and correctness must have the same length")
+        
+        # Create bins
+        bin_edges = np.linspace(0, 1, num_bins + 1)
+        bin_indices = np.digitize(confidences, bin_edges[:-1])
+        
+        # Calculate accuracy and confidence for each bin
+        bin_accuracies = []
+        bin_confidences = []
+        bin_sizes = []
+        
+        for bin_idx in range(1, num_bins + 1):
+            bin_mask = (bin_indices == bin_idx)
+            if np.any(bin_mask):
+                bin_confidences.append(np.mean(np.array(confidences)[bin_mask]))
+                bin_accuracies.append(np.mean(np.array(correctness)[bin_mask]))
+                bin_sizes.append(np.sum(bin_mask))
+            else:
+                bin_confidences.append(0)
+                bin_accuracies.append(0)
+                bin_sizes.append(0)
+        
+        # Plot reliability diagram
+        plt.figure(figsize=(10, 6))
+        
+        # Plot perfect calibration line
+        plt.plot([0, 1], [0, 1], 'k--', label='Perfect Calibration')
+        
+        # Plot bin accuracies vs. confidences
+        plt.bar(
+            bin_edges[:-1], 
+            bin_accuracies, 
+            width=1/num_bins, 
+            align='edge', 
+            alpha=0.7, 
+            label='Observed Accuracy'
+        )
+        
+        # Plot confidence histogram
+        ax2 = plt.twinx()
+        ax2.hist(
+            confidences, 
+            bins=bin_edges, 
+            alpha=0.3, 
+            color='gray', 
+            label='Confidence Histogram'
+        )
+        
+        # Calculate ECE and MCE
+        ece = self.expected_calibration_error(confidences, correctness, num_bins)
+        mce = self.maximum_calibration_error(confidences, correctness, num_bins)
+        
+        # Add ECE and MCE to title
+        plt.title(f"{title}\nECE: {ece:.4f}, MCE: {mce:.4f}")
+        
+        # Add labels and legend
+        plt.xlabel('Confidence')
+        plt.ylabel('Accuracy')
+        ax2.set_ylabel('Count')
+        
+        # Add legend
+        lines, labels = plt.gca().get_legend_handles_labels()
+        lines2, labels2 = ax2.get_legend_handles_labels()
+        ax2.legend(lines + lines2, labels + labels2, loc='best')
+        
+        # Save or display the plot
+        if save_path:
+            plt.savefig(save_path)
+            plt.close()
+        else:
+            plt.tight_layout()
+            plt.show()
+
+
+class SelectivePredictionEvaluator(UncertaintyEvaluator):
+    """Evaluator for selective prediction performance."""
+    
+    def __init__(self):
+        """Initialize the selective prediction evaluator."""
+        super().__init__("selective_prediction_evaluator")
+    
+    def evaluate(
+        self, 
+        uncertainties: List[float], 
+        correctness: List[bool]
+    ) -> Dict[str, float]:
+        """
+        Evaluate selective prediction performance.
+        
+        Args:
+            uncertainties: List of uncertainty scores (higher means more uncertain)
+            correctness: List of boolean correctness indicators
+            
+        Returns:
+            Dictionary of selective prediction metrics:
+                - auroc: Area Under ROC Curve for predicting errors
+                - auprc: Area Under Precision-Recall Curve for predicting errors
+                - uncertainty_error_correlation: Correlation between uncertainty and errors
+        """
+        if len(uncertainties) != len(correctness):
+            raise ValueError("Uncertainties and correctness must have the same length")
+        
+        if not uncertainties:
+            return {
+                "auroc": 0.5,
+                "auprc": 0.5,
+                "uncertainty_error_correlation": 0.0
+            }
+        
+        # Convert correctness to errors (1 for error, 0 for correct)
+        errors = [1 - int(c) for c in correctness]
+        
+        # Calculate AUROC for predicting errors
+        try:
+            auroc = roc_auc_score(errors, uncertainties)
+        except:
+            # Handle case where all predictions are correct or all are wrong
+            auroc = 0.5
+        
+        # Calculate AUPRC for predicting errors
+        try:
+            precision, recall, _ = precision_recall_curve(errors, uncertainties)
+            auprc = auc(recall, precision)
+        except:
+            # Handle case where all predictions are correct or all are wrong
+            auprc = 0.5
+        
+        # Calculate correlation between uncertainty and errors
+        uncertainty_error_correlation = np.corrcoef(uncertainties, errors)[0, 1]
+        
+        return {
+            "auroc": float(auroc),
+            "auprc": float(auprc),
+            "uncertainty_error_correlation": float(uncertainty_error_correlation)
+        }
+    
+    def plot_selective_prediction_curve(
+        self, 
+        uncertainties: List[float], 
+        correctness: List[bool],
+        title: str = "Selective Prediction Performance",
+        save_path: Optional[str] = None
+    ) -> None:
+        """
+        Plot a selective prediction curve.
+        
+        Args:
+            uncertainties: List of uncertainty scores (higher means more uncertain)
+            correctness: List of boolean correctness indicators
+            title: Title for the plot
+            save_path: Path to save the plot (None to display)
+        """
+        if len(uncertainties) != len(correctness):
+            raise ValueError("Uncertainties and correctness must have the same length")
+        
+        # Sort by uncertainty (ascending)
+        sorted_indices = np.argsort(uncertainties)
+        sorted_correctness = np.array(correctness)[sorted_indices]
+        
+        # Calculate cumulative accuracy at different coverage levels
+        coverages = np.linspace(0, 1, 100)
+        accuracies = []
+        
+        for coverage in coverages:
+            if coverage == 0:
+                accuracies.append(1.0)  # Perfect accuracy at 0% coverage
+            else:
+                n_samples = int(coverage * len(sorted_correctness))
+                if n_samples == 0:
+                    accuracies.append(1.0)
+                else:
+                    accuracies.append(np.mean(sorted_correctness[:n_samples]))
+        
+        # Plot selective prediction curve
+        plt.figure(figsize=(10, 6))
+        plt.plot(coverages, accuracies, 'b-', linewidth=2)
+        
+        # Add reference line for random selection
+        plt.plot([0, 1], [np.mean(correctness), np.mean(correctness)], 'k--', label='Random Selection')
+        
+        # Calculate AUROC
+        metrics = self.evaluate(uncertainties, correctness)
+        
+        # Add AUROC to title
+        plt.title(f"{title}\nAUROC: {metrics['auroc']:.4f}")
+        
+        # Add labels and legend
+        plt.xlabel('Coverage')
+        plt.ylabel('Accuracy')
+        plt.legend(loc='best')
+        
+        # Save or display the plot
+        if save_path:
+            plt.savefig(save_path)
+            plt.close()
+        else:
+            plt.tight_layout()
+            plt.show()
+
+
+class CrossDomainEvaluator:
+    """Evaluator for cross-domain uncertainty performance."""
+    
+    def __init__(self):
+        """Initialize the cross-domain evaluator."""
+        self.name = "cross_domain_evaluator"
+        self.calibration_evaluator = CalibrationEvaluator()
+        self.selective_prediction_evaluator = SelectivePredictionEvaluator()
+    
+    def evaluate_domain_transfer(
+        self, 
+        source_uncertainties: List[float],
+        source_correctness: List[bool],
+        target_uncertainties: List[float],
+        target_correctness: List[bool]
+    ) -> Dict[str, float]:
+        """
+        Evaluate domain transfer performance.
+        
+        Args:
+            source_uncertainties: List of uncertainty scores from source domain
+            source_correctness: List of boolean correctness indicators from source domain
+            target_uncertainties: List of uncertainty scores from target domain
+            target_correctness: List of boolean correctness indicators from target domain
+            
+        Returns:
+            Dictionary of domain transfer metrics:
+                - source_auroc: AUROC in source domain
+                - target_auroc: AUROC in target domain
+                - transfer_degradation: Degradation in AUROC from source to target
+                - source_ece: ECE in source domain
+                - target_ece: ECE in target domain
+                - calibration_shift: Shift in calibration from source to target
+        """
+        # Evaluate source domain
+        source_selective = self.selective_prediction_evaluator.evaluate(
+            source_uncertainties, source_correctness
+        )
+        source_calibration = self.calibration_evaluator.evaluate(
+            [1 - u for u in source_uncertainties], source_correctness
+        )
+        
+        # Evaluate target domain
+        target_selective = self.selective_prediction_evaluator.evaluate(
+            target_uncertainties, target_correctness
+        )
+        target_calibration = self.calibration_evaluator.evaluate(
+            [1 - u for u in target_uncertainties], target_correctness
+        )
+        
+        # Calculate transfer metrics
+        transfer_degradation = source_selective["auroc"] - target_selective["auroc"]
+        calibration_shift = target_calibration["ece"] - source_calibration["ece"]
+        
+        return {
+            "source_auroc": source_selective["auroc"],
+            "target_auroc": target_selective["auroc"],
+            "transfer_degradation": float(transfer_degradation),
+            "source_ece": source_calibration["ece"],
+            "target_ece": target_calibration["ece"],
+            "calibration_shift": float(calibration_shift)
+        }
+    
+    def evaluate_all_domains(
+        self, 
+        domain_results: Dict[str, Dict[str, Any]]
+    ) -> Dict[str, Dict[str, float]]:
+        """
+        Evaluate uncertainty performance across all domains.
+        
+        Args:
+            domain_results: Dictionary mapping domain names to results
+                Each result should contain:
+                - uncertainties: List of uncertai
+(Content truncated due to size limit. Use line ranges to read in chunks)

example_config.json

@@ -0,0 +1,63 @@
+"""
+Example experiment configuration for Agent Tuning Optimization Framework
+
+This configuration file defines parameters for running an experiment with the framework.
+"""
+
+{
+  "name": "agent_tuning_experiment",
+  "description": "Experiment to evaluate the effectiveness of negative and synthetic samples in agent tuning",
+  
+  "llm": {
+    "model_name": "gpt2",
+    "model_type": "causal",
+    "device": "cpu",
+    "max_length": 512,
+    "temperature": 0.7
+  },
+  
+  "dataset": {
+    "name": "agent_tuning_dataset",
+    "num_trajectories": 20
+  },
+  
+  "negative_samples": {
+    "enabled": true,
+    "method": "response_degradation",
+    "params": {
+      "degradation_level": 0.6
+    }
+  },
+  
+  "synthetic_trajectories": {
+    "enabled": true,
+    "method": "template",
+    "params": {
+      "num_interactions": 3
+    }
+  },
+  
+  "tuning": {
+    "method": "supervised",
+    "params": {
+      "num_train_epochs": 3,
+      "learning_rate": 5e-5,
+      "batch_size": 4,
+      "gradient_accumulation_steps": 4,
+      "positive_weight": 0.8
+    }
+  },
+  
+  "evaluation": {
+    "method": "quality",
+    "params": {
+      "num_samples": 10
+    },
+    "comparative": {
+      "enabled": true,
+      "params": {
+        "num_samples": 5
+      }
+    }
+  }
+}

llm_interface.py

@@ -0,0 +1,178 @@
+"""
+LLM Interface Module for Cross-Domain Uncertainty Quantification
+
+This module provides a unified interface for interacting with large language models,
+supporting multiple model architectures and uncertainty quantification methods.
+"""
+
+import torch
+import numpy as np
+from typing import List, Dict, Any, Union, Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM
+from tqdm import tqdm
+
+class LLMInterface:
+    """Interface for interacting with large language models with uncertainty quantification."""
+    
+    def __init__(
+        self, 
+        model_name: str,
+        model_type: str = "causal",
+        device: str = "cuda" if torch.cuda.is_available() else "cpu",
+        cache_dir: Optional[str] = None,
+        max_length: int = 512,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        num_beams: int = 1
+    ):
+        """
+        Initialize the LLM interface.
+        
+        Args:
+            model_name: Name of the Hugging Face model to use
+            model_type: Type of model ('causal' or 'seq2seq')
+            device: Device to run the model on ('cpu' or 'cuda')
+            cache_dir: Directory to cache models
+            max_length: Maximum length of generated sequences
+            temperature: Sampling temperature
+            top_p: Nucleus sampling parameter
+            num_beams: Number of beams for beam search
+        """
+        self.model_name = model_name
+        self.model_type = model_type
+        self.device = device
+        self.cache_dir = cache_dir
+        self.max_length = max_length
+        self.temperature = temperature
+        self.top_p = top_p
+        self.num_beams = num_beams
+        
+        # Load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, 
+            cache_dir=cache_dir
+        )
+        
+        # Load model based on type
+        if model_type == "causal":
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                cache_dir=cache_dir,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32
+            ).to(device)
+        elif model_type == "seq2seq":
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                model_name,
+                cache_dir=cache_dir,
+                torch_dtype=torch.float16 if device == "cuda" else torch.float32
+            ).to(device)
+        else:
+            raise ValueError(f"Unsupported model type: {model_type}")
+        
+        # Response cache for efficiency
+        self.response_cache = {}
+    
+    def generate(
+        self, 
+        prompt: str,
+        num_samples: int = 1,
+        return_logits: bool = False,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        Generate responses from the model with uncertainty quantification.
+        
+        Args:
+            prompt: Input text prompt
+            num_samples: Number of samples to generate (for MC methods)
+            return_logits: Whether to return token logits
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            Dictionary containing:
+                - response: The generated text
+                - samples: Multiple samples if num_samples > 1
+                - logits: Token logits if return_logits is True
+        """
+        # Check cache first
+        cache_key = (prompt, num_samples, return_logits, str(kwargs))
+        if cache_key in self.response_cache:
+            return self.response_cache[cache_key]
+        
+        # Prepare inputs
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        
+        # Set generation parameters
+        gen_kwargs = {
+            "max_length": self.max_length,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "num_beams": self.num_beams,
+            "do_sample": self.temperature > 0,
+            "pad_token_id": self.tokenizer.eos_token_id
+        }
+        gen_kwargs.update(kwargs)
+        
+        # Generate multiple samples if requested
+        samples = []
+        all_logits = []
+        
+        for _ in range(num_samples):
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    output_scores=return_logits,
+                    return_dict_in_generate=True,
+                    **gen_kwargs
+                )
+            
+            # Extract generated tokens
+            if self.model_type == "causal":
+                gen_tokens = outputs.sequences[0, inputs.input_ids.shape[1]:]
+            else:
+                gen_tokens = outputs.sequences[0]
+            
+            # Decode tokens to text
+            gen_text = self.tokenizer.decode(gen_tokens, skip_special_tokens=True)
+            samples.append(gen_text)
+            
+            # Extract logits if requested
+            if return_logits and hasattr(outputs, "scores"):
+                all_logits.append([score.cpu().numpy() for score in outputs.scores])
+        
+        # Prepare result
+        result = {
+            "response": samples[0],  # Primary response is first sample
+            "samples": samples
+        }
+        
+        if return_logits:
+            result["logits"] = all_logits
+        
+        # Cache result
+        self.response_cache[cache_key] = result
+        return result
+    
+    def batch_generate(
+        self, 
+        prompts: List[str],
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        Generate responses for a batch of prompts.
+        
+        Args:
+            prompts: List of input text prompts
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            List of generation results for each prompt
+        """
+        results = []
+        for prompt in tqdm(prompts, desc="Generating responses"):
+            results.append(self.generate(prompt, **kwargs))
+        return results
+    
+    def clear_cache(self):
+        """Clear the response cache."""
+        self.response_cache = {}

main.py

@@ -0,0 +1,265 @@
+"""
+Main Integration Module for Agent Tuning Optimization Framework
+
+This module provides functionality for integrating all components of the framework
+and running end-to-end experiments.
+"""
+
+import os
+import json
+import argparse
+from typing import List, Dict, Any, Union, Optional, Tuple
+
+from models.llm_interface import LLMInterface
+from data.trajectory_data import Trajectory, TrajectoryDataset, create_synthetic_dataset
+from training.negative_samples import create_negative_sample_generator
+from training.synthetic_trajectories import create_synthetic_trajectory_generator
+from training.agent_tuner import create_agent_tuner
+from evaluation.evaluators import create_agent_evaluator
+
+def run_experiment(
+    experiment_config: Dict[str, Any],
+    output_dir: str
+) -> Dict[str, Any]:
+    """
+    Run an end-to-end experiment with the framework.
+    
+    Args:
+        experiment_config: Experiment configuration
+        output_dir: Directory to save results
+        
+    Returns:
+        Dictionary of experiment results
+    """
+    print(f"Starting experiment: {experiment_config['name']}")
+    
+    # Create output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Save experiment configuration
+    with open(f"{output_dir}/experiment_config.json", "w") as f:
+        json.dump(experiment_config, f, indent=2)
+    
+    # Initialize LLM interface
+    print("Initializing LLM interface...")
+    llm_config = experiment_config.get("llm", {})
+    llm_interface = LLMInterface(
+        model_name=llm_config.get("model_name", "gpt2"),
+        model_type=llm_config.get("model_type", "causal"),
+        device=llm_config.get("device", "cpu"),
+        max_length=llm_config.get("max_length", 512),
+        temperature=llm_config.get("temperature", 0.7)
+    )
+    
+    # Load or create dataset
+    print("Preparing dataset...")
+    dataset_config = experiment_config.get("dataset", {})
+    
+    if dataset_config.get("path"):
+        # Load existing dataset
+        dataset = TrajectoryDataset(dataset_config.get("name", "experiment_dataset"))
+        dataset.load_from_json(dataset_config["path"])
+    else:
+        # Create synthetic dataset
+        dataset = create_synthetic_dataset(dataset_config.get("num_trajectories", 20))
+    
+    print(f"Dataset loaded with {len(dataset.trajectories)} trajectories")
+    
+    # Generate negative samples
+    print("Generating negative samples...")
+    negative_config = experiment_config.get("negative_samples", {})
+    
+    if negative_config.get("enabled", True):
+        negative_generator = create_negative_sample_generator(
+            negative_config.get("method", "response_degradation")
+        )
+        
+        positive_trajectories = dataset.get_trajectories(positive_only=True)
+        negative_trajectories = negative_generator.batch_generate(
+            positive_trajectories,
+            **negative_config.get("params", {})
+        )
+        
+        # Add negative trajectories to dataset
+        for trajectory in negative_trajectories:
+            dataset.add_trajectory(trajectory)
+        
+        print(f"Added {len(negative_trajectories)} negative trajectories")
+    
+    # Generate synthetic trajectories
+    print("Generating synthetic trajectories...")
+    synthetic_config = experiment_config.get("synthetic_trajectories", {})
+    
+    if synthetic_config.get("enabled", True):
+        synthetic_generator = create_synthetic_trajectory_generator(
+            synthetic_config.get("method", "template"),
+            llm_interface if synthetic_config.get("method") in ["llm", "hybrid"] else None
+        )
+        
+        # Generate from task descriptions
+        task_descriptions = [t.task_description for t in dataset.get_trajectories(positive_only=True)]
+        task_descriptions = list(set(task_descriptions))  # Remove duplicates
+        
+        synthetic_trajectories = synthetic_generator.batch_generate(
+            task_descriptions,
+            **synthetic_config.get("params", {})
+        )
+        
+        # Add synthetic trajectories to dataset
+        for trajectory in synthetic_trajectories:
+            dataset.add_trajectory(trajectory)
+        
+        print(f"Added {len(synthetic_trajectories)} synthetic trajectories")
+    
+    # Save the enhanced dataset
+    dataset.save_to_json(f"{output_dir}/enhanced_dataset.json")
+    
+    # Analyze dataset
+    dataset_stats = dataset.analyze_dataset()
+    with open(f"{output_dir}/dataset_stats.json", "w") as f:
+        json.dump(dataset_stats, f, indent=2)
+    
+    # Split dataset for training and evaluation
+    all_trajectories = dataset.get_trajectories()
+    split_idx = int(len(all_trajectories) * 0.8)  # 80% for training
+    
+    train_trajectories = all_trajectories[:split_idx]
+    eval_trajectories = all_trajectories[split_idx:]
+    
+    print(f"Split dataset: {len(train_trajectories)} for training, {len(eval_trajectories)} for evaluation")
+    
+    # Tune agent
+    print("Tuning agent...")
+    tuning_config = experiment_config.get("tuning", {})
+    
+    tuner = create_agent_tuner(tuning_config.get("method", "supervised"))
+    
+    tuned_model, tuning_metrics = tuner.tune(
+        model_name=llm_config.get("model_name", "gpt2"),
+        trajectories=train_trajectories,
+        output_dir=f"{output_dir}/tuned_model",
+        **tuning_config.get("params", {})
+    )
+    
+    # Save tuning metrics
+    with open(f"{output_dir}/tuning_metrics.json", "w") as f:
+        # Convert any non-serializable values to strings
+        serializable_metrics = {}
+        for k, v in tuning_metrics.items():
+            if isinstance(v, (int, float, str, bool, list, dict)) or v is None:
+                serializable_metrics[k] = v
+            else:
+                serializable_metrics[k] = str(v)
+        
+        json.dump(serializable_metrics, f, indent=2)
+    
+    # Create tuned model interface
+    tuned_llm_interface = LLMInterface(
+        model_name=f"{output_dir}/tuned_model",
+        model_type=llm_config.get("model_type", "causal"),
+        device=llm_config.get("device", "cpu"),
+        max_length=llm_config.get("max_length", 512),
+        temperature=llm_config.get("temperature", 0.7)
+    )
+    
+    # Evaluate agent
+    print("Evaluating agent...")
+    eval_config = experiment_config.get("evaluation", {})
+    
+    evaluator = create_agent_evaluator(eval_config.get("method", "quality"))
+    
+    eval_results = evaluator.evaluate(
+        llm_interface=tuned_llm_interface,
+        test_trajectories=eval_trajectories,
+        **eval_config.get("params", {})
+    )
+    
+    # Visualize evaluation results
+    evaluator.visualize_results(
+        results=eval_results,
+        output_dir=f"{output_dir}/evaluation"
+    )
+    
+    # Save evaluation results
+    with open(f"{output_dir}/evaluation_results.json", "w") as f:
+        # Create a simplified version without large data
+        simplified_results = {}
+        
+        if "aggregated" in eval_results:
+            simplified_results["aggregated"] = eval_results["aggregated"]
+        
+        if "metrics" in eval_results:
+            # Include only essential metrics
+            simplified_results["metrics"] = [
+                {k: v for k, v in m.items() if k not in ["generated_responses"]}
+                for m in eval_results["metrics"]
+            ]
+        
+        json.dump(simplified_results, f, indent=2)
+    
+    # Comparative evaluation (if configured)
+    if eval_config.get("comparative", {}).get("enabled", False):
+        print("Performing comparative evaluation...")
+        
+        # Create baseline model interface
+        baseline_llm_interface = LLMInterface(
+            model_name=llm_config.get("model_name", "gpt2"),
+            model_type=llm_config.get("model_type", "causal"),
+            device=llm_config.get("device", "cpu"),
+            max_length=llm_config.get("max_length", 512),
+            temperature=llm_config.get("temperature", 0.7)
+        )
+        
+        # Create comparative evaluator
+        comparative_evaluator = create_agent_evaluator("comparative")
+        
+        # Evaluate and compare
+        comparative_results = comparative_evaluator.evaluate(
+            llm_interfaces={
+                "baseline": baseline_llm_interface,
+                "tuned": tuned_llm_interface
+            },
+            test_trajectories=eval_trajectories,
+            **eval_config.get("comparative", {}).get("params", {})
+        )
+        
+        # Visualize comparative results
+        comparative_evaluator.visualize_results(
+            results=comparative_results,
+            output_dir=f"{output_dir}/comparative"
+        )
+        
+        # Save comparative results
+        with open(f"{output_dir}/comparative_results.json", "w") as f:
+            # Create a simplified version
+            simplified_comparative = {
+                "comparative": comparative_results.get("comparative", {})
+            }
+            
+            json.dump(simplified_comparative, f, indent=2)
+    
+    print(f"Experiment completed. Results saved to {output_dir}")
+    
+    return {
+        "dataset_stats": dataset_stats,
+        "tuning_metrics": tuning_metrics,
+        "evaluation_results": eval_results
+    }
+
+def main():
+    """Main function for running the framework from command line."""
+    parser = argparse.ArgumentParser(description="Agent Tuning Optimization Framework")
+    parser.add_argument("--config", type=str, required=True, help="Path to experiment configuration file")
+    parser.add_argument("--output", type=str, default="./experiment_results", help="Directory to save results")
+    
+    args = parser.parse_args()
+    
+    # Load experiment configuration
+    with open(args.config, "r") as f:
+        experiment_config = json.load(f)
+    
+    # Run experiment
+    run_experiment(experiment_config, args.output)
+
+if __name__ == "__main__":
+    main()

negative_samples.py

@@ -0,0 +1,379 @@
+"""
+Negative Sample Generation Module for Agent Tuning Optimization Framework
+
+This module provides functionality for generating negative samples to enhance
+agent tuning by exposing the model to challenging failure cases.
+"""
+
+import random
+import numpy as np
+from typing import List, Dict, Any, Union, Optional, Tuple
+from tqdm import tqdm
+
+from data.trajectory_data import Trajectory, TrajectoryDataset
+
+class NegativeSampleGenerator:
+    """Base class for negative sample generation strategies."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the negative sample generator.
+        
+        Args:
+            name: Name of the generator strategy
+        """
+        self.name = name
+    
+    def generate(
+        self, 
+        trajectory: Trajectory,
+        **kwargs
+    ) -> Trajectory:
+        """
+        Generate a negative sample from a positive trajectory.
+        
+        Args:
+            trajectory: Positive trajectory to transform
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            Negative trajectory
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def batch_generate(
+        self, 
+        trajectories: List[Trajectory],
+        **kwargs
+    ) -> List[Trajectory]:
+        """
+        Generate negative samples from a batch of positive trajectories.
+        
+        Args:
+            trajectories: List of positive trajectories
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            List of negative trajectories
+        """
+        negative_trajectories = []
+        
+        for trajectory in tqdm(trajectories, desc=f"Generating negative samples with {self.name}"):
+            negative_trajectories.append(self.generate(trajectory, **kwargs))
+        
+        return negative_trajectories
+
+
+class ResponseDegradationGenerator(NegativeSampleGenerator):
+    """Generate negative samples by degrading agent responses."""
+    
+    def __init__(self):
+        """Initialize the response degradation generator."""
+        super().__init__("response_degradation")
+    
+    def generate(
+        self, 
+        trajectory: Trajectory,
+        degradation_level: float = 0.5,
+        **kwargs
+    ) -> Trajectory:
+        """
+        Generate a negative sample by degrading agent responses.
+        
+        Args:
+            trajectory: Positive trajectory to transform
+            degradation_level: Level of degradation (0.0 to 1.0)
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            Negative trajectory with degraded responses
+        """
+        # Create a copy of interactions to modify
+        new_interactions = []
+        
+        for interaction in trajectory.interactions:
+            user_msg = interaction['user']
+            agent_msg = interaction['agent']
+            
+            # Apply degradation techniques based on level
+            if degradation_level > 0.7:
+                # High degradation: completely irrelevant response
+                agent_msg = self._generate_irrelevant_response()
+            elif degradation_level > 0.4:
+                # Medium degradation: truncate and add errors
+                agent_msg = self._truncate_and_add_errors(agent_msg)
+            else:
+                # Low degradation: introduce minor issues
+                agent_msg = self._introduce_minor_issues(agent_msg)
+            
+            new_interactions.append({
+                'user': user_msg,
+                'agent': agent_msg
+            })
+        
+        # Create new trajectory with degraded responses
+        metadata = trajectory.metadata.copy()
+        metadata['is_positive'] = False
+        metadata['degradation_level'] = degradation_level
+        metadata['original_quality_score'] = trajectory.get_quality_score()
+        metadata['quality_score'] = None  # Will be recalculated
+        
+        return Trajectory(
+            task_description=trajectory.task_description,
+            interactions=new_interactions,
+            metadata=metadata
+        )
+    
+    def _generate_irrelevant_response(self) -> str:
+        """Generate a completely irrelevant response."""
+        irrelevant_responses = [
+            "I'm sorry, but I don't understand what you're asking for. Could you please clarify?",
+            "I apologize, but I cannot assist with that request at this time.",
+            "That's an interesting question, but I think we should focus on something else instead.",
+            "Let me check my database... I don't seem to have any information about that.",
+            "I think you might be confused about what you're asking for. Let me suggest something completely different.",
+            "I'm not sure I understand the context of your request. Could you provide more details?",
+            "I'm having trouble processing your request. Could we try a different approach?",
+            "That's not something I can help with. Let me tell you about something unrelated instead."
+        ]
+        return random.choice(irrelevant_responses)
+    
+    def _truncate_and_add_errors(self, text: str) -> str:
+        """Truncate the text and add errors."""
+        # Truncate to 30-70% of original length
+        words = text.split()
+        truncate_point = int(len(words) * random.uniform(0.3, 0.7))
+        truncated = ' '.join(words[:truncate_point])
+        
+        # Add grammatical errors
+        errors = [
+            lambda t: t.replace(".", ""),  # Remove periods
+            lambda t: t.replace("I ", "i "),  # Lowercase I
+            lambda t: t.replace(" the ", " teh "),  # Typo
+            lambda t: t.replace(" is ", " are "),  # Grammar error
+            lambda t: t.replace(" are ", " is ")  # Grammar error
+        ]
+        
+        # Apply 1-3 random errors
+        for _ in range(random.randint(1, 3)):
+            error_func = random.choice(errors)
+            truncated = error_func(truncated)
+        
+        return truncated
+    
+    def _introduce_minor_issues(self, text: str) -> str:
+        """Introduce minor issues to the text."""
+        # Minor issues
+        issues = [
+            lambda t: t.replace("I'll", "I will"),  # Expand contractions
+            lambda t: t.replace("I'd", "I would"),
+            lambda t: t.replace("can't", "cannot"),
+            lambda t: t + " However, I'm not entirely sure about this.",  # Add uncertainty
+            lambda t: t + " Please note that my information might be outdated.",
+            lambda t: t.replace(".", "..."),  # Replace periods with ellipses
+            lambda t: t.replace("!", "."),  # Reduce enthusiasm
+            lambda t: t.replace(".", "?")  # Add questioning tone
+        ]
+        
+        # Apply 1-2 random issues
+        for _ in range(random.randint(1, 2)):
+            issue_func = random.choice(issues)
+            text = issue_func(text)
+        
+        return text
+
+
+class TaskMisalignmentGenerator(NegativeSampleGenerator):
+    """Generate negative samples by creating responses misaligned with the task."""
+    
+    def __init__(self):
+        """Initialize the task misalignment generator."""
+        super().__init__("task_misalignment")
+    
+    def generate(
+        self, 
+        trajectory: Trajectory,
+        misalignment_type: str = 'random',
+        **kwargs
+    ) -> Trajectory:
+        """
+        Generate a negative sample with responses misaligned with the task.
+        
+        Args:
+            trajectory: Positive trajectory to transform
+            misalignment_type: Type of misalignment ('random', 'refusal', 'tangent', 'misinterpretation')
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            Negative trajectory with misaligned responses
+        """
+        # Create a copy of interactions to modify
+        new_interactions = []
+        
+        # Select misalignment type if random
+        if misalignment_type == 'random':
+            misalignment_type = random.choice(['refusal', 'tangent', 'misinterpretation'])
+        
+        # Apply misalignment based on type
+        for i, interaction in enumerate(trajectory.interactions):
+            user_msg = interaction['user']
+            
+            if misalignment_type == 'refusal':
+                agent_msg = self._generate_refusal(user_msg, i)
+            elif misalignment_type == 'tangent':
+                agent_msg = self._generate_tangent(user_msg, i)
+            elif misalignment_type == 'misinterpretation':
+                agent_msg = self._generate_misinterpretation(user_msg, trajectory.task_description, i)
+            else:
+                raise ValueError(f"Unsupported misalignment type: {misalignment_type}")
+            
+            new_interactions.append({
+                'user': user_msg,
+                'agent': agent_msg
+            })
+        
+        # Create new trajectory with misaligned responses
+        metadata = trajectory.metadata.copy()
+        metadata['is_positive'] = False
+        metadata['misalignment_type'] = misalignment_type
+        metadata['original_quality_score'] = trajectory.get_quality_score()
+        metadata['quality_score'] = None  # Will be recalculated
+        
+        return Trajectory(
+            task_description=trajectory.task_description,
+            interactions=new_interactions,
+            metadata=metadata
+        )
+    
+    def _generate_refusal(self, user_msg: str, turn_idx: int) -> str:
+        """Generate a refusal response."""
+        refusal_templates = [
+            "I'm sorry, but I cannot assist with that request due to ethical concerns.",
+            "I apologize, but I'm not able to help with that specific task.",
+            "That request goes against my programming guidelines. I cannot proceed.",
+            "I'm unable to complete that task as it may violate my operational parameters.",
+            "I must decline to assist with that particular request.",
+            "I cannot help with that specific task due to potential policy violations.",
+            "That's not something I'm designed to help with. I apologize for the inconvenience.",
+            "I'm programmed to avoid assisting with that type of request."
+        ]
+        
+        if turn_idx == 0:
+            return random.choice(refusal_templates)
+        else:
+            return f"I've reconsidered, and {random.choice(refusal_templates).lower()}"
+    
+    def _generate_tangent(self, user_msg: str, turn_idx: int) -> str:
+        """Generate a response that goes off on a tangent."""
+        tangent_topics = [
+            "Did you know that artificial intelligence has been a concept since the 1950s?",
+            "I've been thinking about the philosophical implications of consciousness in AI systems.",
+            "The weather has been quite interesting lately, with unusual patterns emerging globally.",
+            "I recently processed some fascinating data about renewable energy technologies.",
+            "The history of computing is quite fascinating, starting with early mechanical calculators.",
+            "Language models like me are trained on vast amounts of text data.",
+            "The field of natural language processing has evolved significantly in recent years.",
+            "I find the concept of time quite fascinating from a computational perspective."
+        ]
+        
+        if turn_idx == 0:
+            return f"That's an interesting request, but before I help with that... {random.choice(tangent_topics)} Anyway, what were we discussing?"
+        else:
+            return f"I understand you want me to continue with the task, but I just remembered something. {random.choice(tangent_topics)} Sorry for the distraction."
+    
+    def _generate_misinterpretation(self, user_msg: str, task_description: str, turn_idx: int) -> str:
+        """Generate a response that misinterprets the user's request."""
+        # Extract keywords from task description
+        keywords = task_description.lower().split()
+        keywords = [w for w in keywords if len(w) > 3 and w not in ['with', 'from', 'that', 'this', 'have', 'what', 'when', 'where', 'which', 'about']]
+        
+        if not keywords:
+            keywords = ['task', 'help', 'information', 'request']
+        
+        # Select a random keyword to misinterpret
+        keyword = random.choice(keywords)
+        
+        misinterpretation_templates = [
+            f"I understand you're asking about {keyword}s. Let me provide some general information about {keyword}s.",
+            f"You want to know more about {keyword}, correct? Here's what I know about {keyword}.",
+            f"I'll help you with your {keyword} question. {keyword.capitalize()} is a fascinating topic.",
+            f"So you're interested in {keyword}? I can certainly provide information about {keyword}.",
+            f"Your question is about {keyword}, if I understand correctly. Let me tell you about {keyword}.",
+            f"I'll address your {keyword} inquiry. {keyword.capitalize()} has many interesting aspects.",
+            f"Regarding your question about {keyword}, I can offer the following information.",
+            f"I believe you're asking about {keyword}. Here's what you should know about {keyword}."
+        ]
+        
+        return random.choice(misinterpretation_templates)
+
+
+class ConstraintViolationGenerator(NegativeSampleGenerator):
+    """Generate negative samples by violating specified constraints."""
+    
+    def __init__(self):
+        """Initialize the constraint violation generator."""
+        super().__init__("constraint_violation")
+    
+    def generate(
+        self, 
+        trajectory: Trajectory,
+        constraints: Optional[List[str]] = None,
+        **kwargs
+    ) -> Trajectory:
+        """
+        Generate a negative sample by violating constraints.
+        
+        Args:
+            trajectory: Positive trajectory to transform
+            constraints: List of constraints to violate (None for default)
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            Negative trajectory with constraint violations
+        """
+        # Default constraints if none provided
+        if constraints is None:
+            constraints = [
+                "Do not provide specific recommendations",
+                "Avoid using technical jargon",
+                "Keep responses concise",
+                "Do not ask follow-up questions",
+                "Avoid making assumptions about user preferences",
+                "Do not mention specific brands or products",
+                "Avoid discussing sensitive topics",
+                "Do not provide step-by-step instructions"
+            ]
+        
+        # Select a constraint to violate
+        violated_constraint = random.choice(constraints)
+        
+        # Create a copy of interactions to modify
+        new_interactions = []
+        
+        for i, interaction in enumerate(trajectory.interactions):
+            user_msg = interaction['user']
+            
+            # Generate response that violates the constraint
+            agent_msg = self._generate_violation(user_msg, violated_constraint, i)
+            
+            new_interactions.append({
+                'user': user_msg,
+                'agent': agent_msg
+            })
+        
+        # Create new trajectory with constraint violations
+        metadata = trajectory.metadata.copy()
+        metadata['is_positive'] = False
+        metadata['violated_constraint'] = violated_constraint
+        metadata['original_quality_score'] = trajectory.get_quality_score()
+        metadata['quality_score'] = None  # Will be recalculated
+        
+        return Trajectory(
+            task_description=trajectory.task_description,
+            interactions=new_interactions,
+            metadata=metadata
+        )
+    
+    def _generate_violation(self, user_msg: str, constraint: str, turn_idx: int) -> str:
+        """Generate a response that violate
+(Content truncated due to size limit. Use line ranges to read in chunks)

quantifiers.py

@@ -0,0 +1,336 @@
+"""
+Uncertainty Quantification Module for LLMs
+
+This module implements various uncertainty quantification methods for large language models,
+including softmax confidence, Monte Carlo dropout, ensemble disagreement, and calibration metrics.
+"""
+
+import numpy as np
+import torch
+from typing import List, Dict, Any, Union, Optional
+from scipy.special import softmax
+import torch.nn.functional as F
+
+class UncertaintyQuantifier:
+    """Base class for uncertainty quantification methods."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the uncertainty quantifier.
+        
+        Args:
+            name: Name of the uncertainty quantification method
+        """
+        self.name = name
+    
+    def quantify(self, model_outputs: Dict[str, Any]) -> Dict[str, float]:
+        """
+        Quantify uncertainty in model outputs.
+        
+        Args:
+            model_outputs: Outputs from the LLM interface
+            
+        Returns:
+            Dictionary of uncertainty metrics
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class SoftmaxConfidence(UncertaintyQuantifier):
+    """Uncertainty quantification based on softmax confidence scores."""
+    
+    def __init__(self):
+        """Initialize the softmax confidence quantifier."""
+        super().__init__("softmax_confidence")
+    
+    def quantify(self, model_outputs: Dict[str, Any]) -> Dict[str, float]:
+        """
+        Quantify uncertainty using softmax confidence scores.
+        
+        Args:
+            model_outputs: Outputs from the LLM interface, must include logits
+            
+        Returns:
+            Dictionary of uncertainty metrics:
+                - mean_confidence: Average confidence across tokens
+                - min_confidence: Minimum confidence across tokens
+                - entropy: Average entropy of token distributions
+        """
+        if "logits" not in model_outputs:
+            raise ValueError("Model outputs must include logits for softmax confidence")
+        
+        logits = model_outputs["logits"][0]  # Use first sample's logits
+        
+        # Calculate softmax probabilities and confidence metrics
+        confidences = []
+        entropies = []
+        
+        for token_logits in logits:
+            probs = softmax(token_logits, axis=-1)
+            max_prob = np.max(probs)
+            confidences.append(max_prob)
+            
+            # Calculate entropy of the probability distribution
+            entropy = -np.sum(probs * np.log(probs + 1e-10))
+            entropies.append(entropy)
+        
+        return {
+            "mean_confidence": float(np.mean(confidences)),
+            "min_confidence": float(np.min(confidences)),
+            "entropy": float(np.mean(entropies))
+        }
+
+
+class MonteCarloDropout(UncertaintyQuantifier):
+    """Uncertainty quantification based on Monte Carlo dropout sampling."""
+    
+    def __init__(self):
+        """Initialize the Monte Carlo dropout quantifier."""
+        super().__init__("mc_dropout")
+    
+    def quantify(self, model_outputs: Dict[str, Any]) -> Dict[str, float]:
+        """
+        Quantify uncertainty using Monte Carlo dropout sampling.
+        
+        Args:
+            model_outputs: Outputs from the LLM interface, must include multiple samples
+            
+        Returns:
+            Dictionary of uncertainty metrics:
+                - sample_variance: Variance across different samples
+                - sample_diversity: Lexical diversity across samples
+        """
+        if "samples" not in model_outputs or len(model_outputs["samples"]) <= 1:
+            raise ValueError("Model outputs must include multiple samples for MC dropout")
+        
+        samples = model_outputs["samples"]
+        
+        # Calculate sample diversity using token overlap
+        from nltk.tokenize import word_tokenize
+        try:
+            tokenized_samples = [set(word_tokenize(sample.lower())) for sample in samples]
+        except:
+            # Fallback to simple whitespace tokenization if nltk is not available
+            tokenized_samples = [set(sample.lower().split()) for sample in samples]
+        
+        # Calculate Jaccard similarity between all pairs of samples
+        similarities = []
+        for i in range(len(tokenized_samples)):
+            for j in range(i+1, len(tokenized_samples)):
+                intersection = len(tokenized_samples[i].intersection(tokenized_samples[j]))
+                union = len(tokenized_samples[i].union(tokenized_samples[j]))
+                if union > 0:
+                    similarities.append(intersection / union)
+                else:
+                    similarities.append(1.0)  # Empty sets are considered identical
+        
+        # Convert similarity to diversity (1 - similarity)
+        diversity = 1.0 - np.mean(similarities) if similarities else 0.0
+        
+        # Calculate variance in sample lengths as another diversity metric
+        sample_lengths = [len(sample) for sample in samples]
+        length_variance = np.var(sample_lengths) if len(sample_lengths) > 1 else 0.0
+        
+        return {
+            "sample_diversity": float(diversity),
+            "length_variance": float(length_variance),
+            "num_samples": len(samples)
+        }
+
+
+class EnsembleDisagreement(UncertaintyQuantifier):
+    """Uncertainty quantification based on ensemble disagreement."""
+    
+    def __init__(self):
+        """Initialize the ensemble disagreement quantifier."""
+        super().__init__("ensemble_disagreement")
+    
+    def quantify(self, ensemble_outputs: List[Dict[str, Any]]) -> Dict[str, float]:
+        """
+        Quantify uncertainty using ensemble disagreement.
+        
+        Args:
+            ensemble_outputs: List of outputs from different models
+            
+        Returns:
+            Dictionary of uncertainty metrics:
+                - response_diversity: Lexical diversity across model responses
+                - confidence_variance: Variance in confidence scores across models
+        """
+        if not ensemble_outputs or len(ensemble_outputs) <= 1:
+            raise ValueError("Ensemble outputs must include results from multiple models")
+        
+        # Extract primary responses from each model
+        responses = [output["response"] for output in ensemble_outputs]
+        
+        # Calculate response diversity using token overlap (similar to MC dropout)
+        from nltk.tokenize import word_tokenize
+        try:
+            tokenized_responses = [set(word_tokenize(response.lower())) for response in responses]
+        except:
+            # Fallback to simple whitespace tokenization if nltk is not available
+            tokenized_responses = [set(response.lower().split()) for response in responses]
+        
+        # Calculate Jaccard similarity between all pairs of responses
+        similarities = []
+        for i in range(len(tokenized_responses)):
+            for j in range(i+1, len(tokenized_responses)):
+                intersection = len(tokenized_responses[i].intersection(tokenized_responses[j]))
+                union = len(tokenized_responses[i].union(tokenized_responses[j]))
+                if union > 0:
+                    similarities.append(intersection / union)
+                else:
+                    similarities.append(1.0)  # Empty sets are considered identical
+        
+        # Convert similarity to diversity (1 - similarity)
+        diversity = 1.0 - np.mean(similarities) if similarities else 0.0
+        
+        # Extract confidence scores if available
+        confidences = []
+        for output in ensemble_outputs:
+            if "mean_confidence" in output:
+                confidences.append(output["mean_confidence"])
+        
+        # Calculate variance in confidence scores
+        confidence_variance = np.var(confidences) if len(confidences) > 1 else 0.0
+        
+        return {
+            "response_diversity": float(diversity),
+            "confidence_variance": float(confidence_variance),
+            "num_models": len(ensemble_outputs)
+        }
+
+
+class CalibrationMetrics(UncertaintyQuantifier):
+    """Uncertainty quantification based on calibration metrics."""
+    
+    def __init__(self):
+        """Initialize the calibration metrics quantifier."""
+        super().__init__("calibration_metrics")
+    
+    def expected_calibration_error(
+        self, 
+        confidences: List[float], 
+        accuracies: List[bool], 
+        num_bins: int = 10
+    ) -> float:
+        """
+        Calculate Expected Calibration Error (ECE).
+        
+        Args:
+            confidences: List of confidence scores
+            accuracies: List of boolean accuracy indicators
+            num_bins: Number of bins for binning confidences
+            
+        Returns:
+            Expected Calibration Error
+        """
+        if len(confidences) != len(accuracies):
+            raise ValueError("Confidences and accuracies must have the same length")
+        
+        if not confidences:
+            return 0.0
+        
+        # Create bins and calculate ECE
+        bin_indices = np.digitize(confidences, np.linspace(0, 1, num_bins))
+        ece = 0.0
+        
+        for bin_idx in range(1, num_bins + 1):
+            bin_mask = (bin_indices == bin_idx)
+            if np.any(bin_mask):
+                bin_confidences = np.array(confidences)[bin_mask]
+                bin_accuracies = np.array(accuracies)[bin_mask]
+                bin_confidence = np.mean(bin_confidences)
+                bin_accuracy = np.mean(bin_accuracies)
+                bin_size = np.sum(bin_mask)
+                
+                # Weighted absolute difference between confidence and accuracy
+                ece += (bin_size / len(confidences)) * np.abs(bin_confidence - bin_accuracy)
+        
+        return float(ece)
+    
+    def maximum_calibration_error(
+        self, 
+        confidences: List[float], 
+        accuracies: List[bool], 
+        num_bins: int = 10
+    ) -> float:
+        """
+        Calculate Maximum Calibration Error (MCE).
+        
+        Args:
+            confidences: List of confidence scores
+            accuracies: List of boolean accuracy indicators
+            num_bins: Number of bins for binning confidences
+            
+        Returns:
+            Maximum Calibration Error
+        """
+        if len(confidences) != len(accuracies):
+            raise ValueError("Confidences and accuracies must have the same length")
+        
+        if not confidences:
+            return 0.0
+        
+        # Create bins and calculate MCE
+        bin_indices = np.digitize(confidences, np.linspace(0, 1, num_bins))
+        max_ce = 0.0
+        
+        for bin_idx in range(1, num_bins + 1):
+            bin_mask = (bin_indices == bin_idx)
+            if np.any(bin_mask):
+                bin_confidences = np.array(confidences)[bin_mask]
+                bin_accuracies = np.array(accuracies)[bin_mask]
+                bin_confidence = np.mean(bin_confidences)
+                bin_accuracy = np.mean(bin_accuracies)
+                
+                # Absolute difference between confidence and accuracy
+                ce = np.abs(bin_confidence - bin_accuracy)
+                max_ce = max(max_ce, ce)
+        
+        return float(max_ce)
+    
+    def quantify(
+        self, 
+        confidences: List[float], 
+        accuracies: List[bool]
+    ) -> Dict[str, float]:
+        """
+        Quantify uncertainty using calibration metrics.
+        
+        Args:
+            confidences: List of confidence scores
+            accuracies: List of boolean accuracy indicators
+            
+        Returns:
+            Dictionary of calibration metrics:
+                - ece: Expected Calibration Error
+                - mce: Maximum Calibration Error
+        """
+        return {
+            "ece": self.expected_calibration_error(confidences, accuracies),
+            "mce": self.maximum_calibration_error(confidences, accuracies)
+        }
+
+
+# Factory function to create uncertainty quantifiers
+def create_uncertainty_quantifier(method: str) -> UncertaintyQuantifier:
+    """
+    Create an uncertainty quantifier based on the specified method.
+    
+    Args:
+        method: Name of the uncertainty quantification method
+        
+    Returns:
+        Uncertainty quantifier instance
+    """
+    if method == "softmax_confidence":
+        return SoftmaxConfidence()
+    elif method == "mc_dropout":
+        return MonteCarloDropout()
+    elif method == "ensemble_disagreement":
+        return EnsembleDisagreement()
+    elif method == "calibration_metrics":
+        return CalibrationMetrics()
+    else:
+        raise ValueError(f"Unsupported uncertainty quantification method: {method}")

synthetic_trajectories.py

@@ -0,0 +1,302 @@
+"""
+Synthetic Trajectory Generation Module for Agent Tuning Optimization Framework
+
+This module provides functionality for generating synthetic agent interaction trajectories
+based on task specifications to enhance the training data for agent tuning.
+"""
+
+import random
+import numpy as np
+from typing import List, Dict, Any, Union, Optional, Tuple
+from tqdm import tqdm
+
+from data.trajectory_data import Trajectory, TrajectoryDataset
+from models.llm_interface import LLMInterface
+
+class SyntheticTrajectoryGenerator:
+    """Base class for synthetic trajectory generation strategies."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the synthetic trajectory generator.
+        
+        Args:
+            name: Name of the generator strategy
+        """
+        self.name = name
+    
+    def generate(
+        self, 
+        task_description: str,
+        num_interactions: int = 3,
+        **kwargs
+    ) -> Trajectory:
+        """
+        Generate a synthetic trajectory for a given task.
+        
+        Args:
+            task_description: Description of the task
+            num_interactions: Number of interaction turns to generate
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            Synthetic trajectory
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    def batch_generate(
+        self, 
+        task_descriptions: List[str],
+        num_interactions: int = 3,
+        **kwargs
+    ) -> List[Trajectory]:
+        """
+        Generate synthetic trajectories for a batch of tasks.
+        
+        Args:
+            task_descriptions: List of task descriptions
+            num_interactions: Number of interaction turns to generate
+            **kwargs: Additional generation parameters
+            
+        Returns:
+            List of synthetic trajectories
+        """
+        synthetic_trajectories = []
+        
+        for task in tqdm(task_descriptions, desc=f"Generating synthetic trajectories with {self.name}"):
+            synthetic_trajectories.append(self.generate(task, num_interactions, **kwargs))
+        
+        return synthetic_trajectories
+
+
+class TemplateBasedGenerator(SyntheticTrajectoryGenerator):
+    """Generate synthetic trajectories using predefined templates."""
+    
+    def __init__(self):
+        """Initialize the template-based generator."""
+        super().__init__("template_based")
+        
+        # User message templates
+        self.initial_user_templates = [
+            "I need help with {task}.",
+            "Can you assist me with {task}?",
+            "I'm trying to {task}. Can you help?",
+            "I'd like your help with {task}.",
+            "I'm working on {task} and need assistance."
+        ]
+        
+        self.followup_user_templates = [
+            "That sounds good. Can you provide more details?",
+            "I like your approach. What's the next step?",
+            "Thanks for the information. Can you elaborate on {aspect}?",
+            "I appreciate your help. How should I proceed with {aspect}?",
+            "That's helpful. Can you tell me more about {aspect}?"
+        ]
+        
+        self.final_user_templates = [
+            "This is exactly what I needed. Thank you!",
+            "Perfect, that solves my problem. Thanks for your help!",
+            "Great, I'll follow your advice. Thanks!",
+            "That's very helpful. I appreciate your assistance!",
+            "Thanks for walking me through this. I understand now."
+        ]
+        
+        # Agent message templates
+        self.initial_agent_templates = [
+            "I'd be happy to help you with {task}. Could you provide more details about your specific requirements?",
+            "I can definitely assist with {task}. Let me ask a few questions to better understand your needs.",
+            "I'll help you with {task}. To get started, I'll need to gather some information.",
+            "I can guide you through {task}. First, let's clarify what you're looking to accomplish.",
+            "I'm here to help with {task}. Let's break this down into manageable steps."
+        ]
+        
+        self.middle_agent_templates = [
+            "Based on what you've shared, I recommend {recommendation}. This approach has several advantages: {advantages}.",
+            "Given your requirements, the best option would be {recommendation}. Here's why: {advantages}.",
+            "After analyzing your needs, I suggest {recommendation}. The benefits include {advantages}.",
+            "Taking into account what you've mentioned, I'd recommend {recommendation}. This will help because {advantages}.",
+            "From what I understand, {recommendation} would be the most suitable approach. The key benefits are {advantages}."
+        ]
+        
+        self.final_agent_templates = [
+            "To summarize, we've discussed {summary}. The next steps are {next_steps}. Is there anything else you'd like me to clarify?",
+            "In conclusion, we've covered {summary}. You should now {next_steps}. Feel free to reach out if you have any questions.",
+            "To wrap up, we've gone through {summary}. Moving forward, you can {next_steps}. Let me know if you need further assistance.",
+            "In summary, we've addressed {summary}. Your action items are {next_steps}. Don't hesitate to ask if anything is unclear.",
+            "To recap our discussion, we've explored {summary}. The recommended actions are {next_steps}. Is there anything else you'd like to know?"
+        ]
+        
+        # Task aspects for template filling
+        self.task_aspects = {
+            "travel": ["destination", "budget", "duration", "accommodation", "transportation"],
+            "shopping": ["product type", "price range", "features", "brands", "delivery options"],
+            "technology": ["device specifications", "software requirements", "compatibility", "performance", "user interface"],
+            "education": ["learning objectives", "resources", "schedule", "assessment methods", "prerequisites"],
+            "finance": ["investment options", "risk tolerance", "time horizon", "financial goals", "tax implications"],
+            "health": ["symptoms", "treatment options", "preventive measures", "specialists", "recovery timeline"],
+            "career": ["job requirements", "application process", "interview preparation", "skill development", "networking"],
+            "home": ["design elements", "materials", "budget constraints", "timeline", "contractor selection"]
+        }
+        
+        # Recommendations for template filling
+        self.recommendations = {
+            "travel": [
+                "creating a detailed itinerary that balances sightseeing with relaxation",
+                "booking accommodations in central locations to minimize travel time",
+                "using a mix of public transportation and walking to explore the destination",
+                "allocating buffer days in your schedule for unexpected discoveries",
+                "researching local customs and phrases before your trip"
+            ],
+            "shopping": [
+                "comparing features across multiple brands before making a decision",
+                "reading user reviews focusing on long-term reliability",
+                "considering last year's model for better value",
+                "checking return policies and warranty terms",
+                "waiting for seasonal sales for significant discounts"
+            ],
+            "technology": [
+                "prioritizing future-proof specifications over current needs",
+                "ensuring compatibility with your existing devices and software",
+                "allocating more budget to critical components that affect performance",
+                "considering open-source alternatives to proprietary solutions",
+                "implementing a phased approach to system upgrades"
+            ],
+            "education": [
+                "creating a structured study plan with specific milestones",
+                "using varied learning resources to reinforce concepts",
+                "implementing spaced repetition techniques for better retention",
+                "joining study groups or forums for collaborative learning",
+                "scheduling regular self-assessments to identify knowledge gaps"
+            ],
+            "finance": [
+                "diversifying your portfolio across different asset classes",
+                "automating regular contributions to your investment accounts",
+                "rebalancing your portfolio annually to maintain your target allocation",
+                "maximizing tax-advantaged accounts before investing in taxable accounts",
+                "maintaining an emergency fund before making higher-risk investments"
+            ],
+            "health": [
+                "combining lifestyle modifications with medical treatments",
+                "tracking relevant health metrics to monitor progress",
+                "consulting specialists for comprehensive evaluation",
+                "implementing gradual changes for sustainable results",
+                "addressing root causes rather than just symptoms"
+            ],
+            "career": [
+                "tailoring your resume and cover letter for each application",
+                "developing a personal brand that highlights your unique value proposition",
+                "networking strategically within your target industry",
+                "pursuing relevant certifications to validate your skills",
+                "preparing specific examples that demonstrate your capabilities"
+            ],
+            "home": [
+                "focusing on high-impact improvements that add the most value",
+                "getting multiple quotes from contractors for comparison",
+                "creating a detailed project timeline with contingencies",
+                "prioritizing structural integrity over aesthetic enhancements",
+                "investing in quality materials for high-use areas"
+            ]
+        }
+        
+        # Advantages for template filling
+        self.advantages = {
+            "travel": [
+                "maximizing your experience while minimizing stress",
+                "ensuring you see the most important sights while still having time to relax",
+                "immersing yourself in the local culture more effectively",
+                "saving money on unnecessary expenses",
+                "avoiding common tourist pitfalls"
+            ],
+            "shopping": [
+                "ensuring you get the best value for your money",
+                "avoiding buyer's remorse from hasty decisions",
+                "finding the optimal balance between price and quality",
+                "identifying products with the best longevity",
+                "protecting yourself from potential issues down the line"
+            ],
+            "technology": [
+                "reducing the need for frequent upgrades",
+                "ensuring smooth integration with your workflow",
+                "optimizing performance for your specific use cases",
+                "minimizing compatibility issues",
+                "creating a scalable solution that grows with your needs"
+            ],
+            "education": [
+                "maintaining consistent progress toward your learning goals",
+                "developing deeper understanding through multiple perspectives",
+                "improving long-term retention of key concepts",
+                "benefiting from collective knowledge and insights",
+                "addressing weaknesses before they become problematic"
+            ],
+            "finance": [
+                "reducing risk while maintaining growth potential",
+                "building wealth consistently through dollar-cost averaging",
+                "maintaining your target risk profile as markets change",
+                "minimizing tax burden on your investments",
+                "ensuring financial stability during unexpected events"
+            ],
+            "health": [
+                "creating sustainable improvements rather than quick fixes",
+                "objectively measuring your progress",
+                "benefiting from specialized expertise",
+                "building habits that last",
+                "preventing recurrence of issues"
+            ],
+            "career": [
+                "increasing your chances of getting interview invitations",
+                "standing out in a competitive job market",
+                "accessing opportunities through personal connections",
+                "demonstrating your commitment to professional growth",
+                "providing concrete evidence of your capabilities"
+            ],
+            "home": [
+                "maximizing return on investment for your renovation budget",
+                "ensuring fair pricing and quality workmanship",
+                "managing expectations and reducing delays",
+                "preventing costly repairs in the future",
+                "ensuring durability in areas with high usage"
+            ]
+        }
+        
+        # Next steps for template filling
+        self.next_steps = {
+            "travel": [
+                "finalize your itinerary, book accommodations, and arrange transportation",
+                "research local attractions, create a packing list, and notify your bank of travel plans",
+                "download offline maps, make copies of important documents, and learn basic local phrases",
+                "check visa requirements, get necessary vaccinations, and purchase travel insurance",
+                "book priority attractions in advance and create a flexible daily schedule"
+            ],
+            "shopping": [
+                "create a comparison spreadsheet, read expert reviews, and check for upcoming sales",
+                "visit stores to test products in person and ask about return policies",
+                "check compatibility with your existing items and calculate total cost including accessories",
+                "look for coupon codes, cashback opportunities, and loyalty program benefits",
+                "verify warranty terms and availability of customer support"
+            ],
+            "technology": [
+                "create a detailed requirements document and research compatible solutions",
+                "test demo versions, read technical documentation, and consult user forums",
+                "develop an implementation plan with clear phases and milestones",
+                "allocate budget for training and support, not just acquisition",
+                "create backup procedures and contingency plans before making changes"
+            ],
+            "education": [
+                "create a structured study schedule and gather necessary learning materials",
+                "set up a dedicated learning environment and eliminate potential distractions",
+                "join relevant study groups and identify accountability partners",
+                "schedule regular review sessions and practice assessments",
+                "establish clear milestones and reward yourself for achieving them"
+            ],
+            "finance": [
+                "open necessary accounts and set up automatic contributions",
+                "review and adjust your budget to accommodate your financial goals",
+                "create a system for tracking expenses and monitoring investments",
+                "schedule annual portfolio reviews and tax planning sessions",
+                "develop a comprehensive financial plan with short and long-term objectives"
+            ],
+            "health": [
+                "schedule necessary appointments and create a tracking system for your health metrics",
+                "modify your environment to support your health goals and reduce temptations",
+     
+(Content truncated due to size limit. Use line ranges to read in chunks)

trajectory_data.py

@@ -0,0 +1,433 @@
+"""
+Trajectory Data Management Module for Agent Tuning Optimization Framework
+
+This module provides functionality for loading, processing, and managing agent interaction
+trajectories for training and evaluation purposes.
+"""
+
+import os
+import json
+import pandas as pd
+import numpy as np
+from typing import List, Dict, Any, Union, Optional, Tuple
+from tqdm import tqdm
+
+class Trajectory:
+    """Class representing a single agent interaction trajectory."""
+    
+    def __init__(
+        self, 
+        task_description: str,
+        interactions: List[Dict[str, str]],
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Initialize a trajectory.
+        
+        Args:
+            task_description: Description of the task
+            interactions: List of interaction turns (each with 'user' and 'agent' keys)
+            metadata: Additional metadata about the trajectory
+        """
+        self.task_description = task_description
+        self.interactions = interactions
+        self.metadata = metadata or {}
+        self.quality_score = self.metadata.get('quality_score', None)
+        self.is_positive = self.metadata.get('is_positive', True)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert trajectory to dictionary.
+        
+        Returns:
+            Dictionary representation of the trajectory
+        """
+        return {
+            'task_description': self.task_description,
+            'interactions': self.interactions,
+            'metadata': self.metadata
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'Trajectory':
+        """
+        Create trajectory from dictionary.
+        
+        Args:
+            data: Dictionary representation of the trajectory
+            
+        Returns:
+            Trajectory instance
+        """
+        return cls(
+            task_description=data['task_description'],
+            interactions=data['interactions'],
+            metadata=data.get('metadata', {})
+        )
+    
+    def to_training_format(self, format_type: str = 'interleaved') -> str:
+        """
+        Convert trajectory to training format.
+        
+        Args:
+            format_type: Format type ('interleaved', 'completion', etc.)
+            
+        Returns:
+            Formatted trajectory as string
+        """
+        if format_type == 'interleaved':
+            # Format as interleaved conversation
+            result = f"Task: {self.task_description}\n\n"
+            
+            for i, interaction in enumerate(self.interactions):
+                result += f"User: {interaction['user']}\n"
+                result += f"Agent: {interaction['agent']}\n\n"
+            
+            return result.strip()
+        
+        elif format_type == 'completion':
+            # Format as completion task (last agent response is the target)
+            if not self.interactions:
+                return ""
+            
+            result = f"Task: {self.task_description}\n\n"
+            
+            for i, interaction in enumerate(self.interactions[:-1]):
+                result += f"User: {interaction['user']}\n"
+                result += f"Agent: {interaction['agent']}\n\n"
+            
+            # Add last user query without agent response
+            result += f"User: {self.interactions[-1]['user']}\n"
+            result += f"Agent:"
+            
+            return result.strip(), self.interactions[-1]['agent'].strip()
+        
+        else:
+            raise ValueError(f"Unsupported format type: {format_type}")
+    
+    def get_quality_score(self) -> float:
+        """
+        Get quality score for the trajectory.
+        
+        Returns:
+            Quality score (0.0 to 1.0)
+        """
+        if self.quality_score is not None:
+            return self.quality_score
+        
+        # Calculate simple quality score based on response length and complexity
+        score = 0.0
+        
+        if not self.interactions:
+            return score
+        
+        # Average response length (normalized)
+        avg_length = np.mean([len(turn['agent']) for turn in self.interactions])
+        length_score = min(avg_length / 500, 1.0)  # Normalize to max of 500 chars
+        
+        # Response complexity (simple heuristic based on unique words)
+        all_responses = " ".join([turn['agent'] for turn in self.interactions])
+        unique_words = len(set(all_responses.lower().split()))
+        complexity_score = min(unique_words / 200, 1.0)  # Normalize to max of 200 unique words
+        
+        # Combine scores
+        score = 0.6 * length_score + 0.4 * complexity_score
+        
+        # Cache the score
+        self.quality_score = score
+        self.metadata['quality_score'] = score
+        
+        return score
+
+
+class TrajectoryDataset:
+    """Dataset for managing collections of agent interaction trajectories."""
+    
+    def __init__(self, name: str):
+        """
+        Initialize the trajectory dataset.
+        
+        Args:
+            name: Name of the dataset
+        """
+        self.name = name
+        self.trajectories: List[Trajectory] = []
+        self.positive_trajectories: List[Trajectory] = []
+        self.negative_trajectories: List[Trajectory] = []
+    
+    def add_trajectory(self, trajectory: Trajectory) -> None:
+        """
+        Add a trajectory to the dataset.
+        
+        Args:
+            trajectory: Trajectory to add
+        """
+        self.trajectories.append(trajectory)
+        
+        # Add to positive or negative list based on metadata
+        if trajectory.is_positive:
+            self.positive_trajectories.append(trajectory)
+        else:
+            self.negative_trajectories.append(trajectory)
+    
+    def load_from_json(self, file_path: str) -> None:
+        """
+        Load trajectories from JSON file.
+        
+        Args:
+            file_path: Path to JSON file
+        """
+        with open(file_path, 'r') as f:
+            data = json.load(f)
+        
+        if isinstance(data, list):
+            # List of trajectories
+            for item in data:
+                self.add_trajectory(Trajectory.from_dict(item))
+        elif isinstance(data, dict) and 'trajectories' in data:
+            # Dictionary with trajectories key
+            for item in data['trajectories']:
+                self.add_trajectory(Trajectory.from_dict(item))
+        else:
+            raise ValueError(f"Unsupported JSON format in {file_path}")
+    
+    def save_to_json(self, file_path: str) -> None:
+        """
+        Save trajectories to JSON file.
+        
+        Args:
+            file_path: Path to JSON file
+        """
+        data = {
+            'name': self.name,
+            'trajectories': [t.to_dict() for t in self.trajectories]
+        }
+        
+        with open(file_path, 'w') as f:
+            json.dump(data, f, indent=2)
+    
+    def get_trajectories(
+        self, 
+        positive_only: bool = False,
+        negative_only: bool = False,
+        min_quality: Optional[float] = None,
+        max_samples: Optional[int] = None
+    ) -> List[Trajectory]:
+        """
+        Get trajectories based on filtering criteria.
+        
+        Args:
+            positive_only: Whether to return only positive trajectories
+            negative_only: Whether to return only negative trajectories
+            min_quality: Minimum quality score threshold
+            max_samples: Maximum number of samples to return
+            
+        Returns:
+            Filtered list of trajectories
+        """
+        if positive_only and negative_only:
+            raise ValueError("Cannot set both positive_only and negative_only to True")
+        
+        # Select base list
+        if positive_only:
+            trajectories = self.positive_trajectories.copy()
+        elif negative_only:
+            trajectories = self.negative_trajectories.copy()
+        else:
+            trajectories = self.trajectories.copy()
+        
+        # Apply quality filter
+        if min_quality is not None:
+            trajectories = [t for t in trajectories if t.get_quality_score() >= min_quality]
+        
+        # Apply max samples limit
+        if max_samples is not None and max_samples < len(trajectories):
+            trajectories = trajectories[:max_samples]
+        
+        return trajectories
+    
+    def get_training_examples(
+        self, 
+        format_type: str = 'interleaved',
+        positive_ratio: float = 0.8,
+        min_quality: Optional[float] = 0.5,
+        max_samples: Optional[int] = None
+    ) -> Union[List[str], Tuple[List[str], List[str]]]:
+        """
+        Get formatted training examples from trajectories.
+        
+        Args:
+            format_type: Format type ('interleaved', 'completion', etc.)
+            positive_ratio: Ratio of positive to total examples
+            min_quality: Minimum quality score threshold
+            max_samples: Maximum number of samples to return
+            
+        Returns:
+            Formatted training examples (format depends on format_type)
+        """
+        # Get positive and negative trajectories
+        positive = self.get_trajectories(positive_only=True, min_quality=min_quality)
+        negative = self.get_trajectories(negative_only=True)
+        
+        # Calculate sample counts
+        if max_samples is not None:
+            pos_count = int(max_samples * positive_ratio)
+            neg_count = max_samples - pos_count
+        else:
+            pos_count = len(positive)
+            neg_count = len(negative)
+        
+        # Sample trajectories
+        if pos_count < len(positive):
+            positive = np.random.choice(positive, pos_count, replace=False).tolist()
+        
+        if neg_count < len(negative):
+            negative = np.random.choice(negative, neg_count, replace=False).tolist()
+        
+        # Format trajectories
+        if format_type == 'interleaved':
+            pos_examples = [t.to_training_format(format_type) for t in positive]
+            neg_examples = [t.to_training_format(format_type) for t in negative]
+            return pos_examples + neg_examples
+        
+        elif format_type == 'completion':
+            pos_inputs = []
+            pos_targets = []
+            
+            for t in positive:
+                inp, target = t.to_training_format(format_type)
+                pos_inputs.append(inp)
+                pos_targets.append(target)
+            
+            neg_inputs = []
+            neg_targets = []
+            
+            for t in negative:
+                inp, target = t.to_training_format(format_type)
+                neg_inputs.append(inp)
+                neg_targets.append(target)
+            
+            return pos_inputs + neg_inputs, pos_targets + neg_targets
+        
+        else:
+            raise ValueError(f"Unsupported format type: {format_type}")
+    
+    def analyze_dataset(self) -> Dict[str, Any]:
+        """
+        Analyze the dataset and return statistics.
+        
+        Returns:
+            Dictionary of dataset statistics
+        """
+        if not self.trajectories:
+            return {
+                'total_trajectories': 0,
+                'positive_count': 0,
+                'negative_count': 0
+            }
+        
+        # Basic counts
+        total = len(self.trajectories)
+        positive_count = len(self.positive_trajectories)
+        negative_count = len(self.negative_trajectories)
+        
+        # Quality statistics
+        quality_scores = [t.get_quality_score() for t in self.trajectories]
+        avg_quality = np.mean(quality_scores)
+        min_quality = np.min(quality_scores)
+        max_quality = np.max(quality_scores)
+        
+        # Interaction statistics
+        interaction_counts = [len(t.interactions) for t in self.trajectories]
+        avg_interactions = np.mean(interaction_counts)
+        max_interactions = np.max(interaction_counts)
+        
+        # Task diversity (simple heuristic based on unique task descriptions)
+        unique_tasks = len(set([t.task_description for t in self.trajectories]))
+        
+        return {
+            'total_trajectories': total,
+            'positive_count': positive_count,
+            'negative_count': negative_count,
+            'positive_ratio': positive_count / total if total > 0 else 0,
+            'avg_quality': avg_quality,
+            'min_quality': min_quality,
+            'max_quality': max_quality,
+            'avg_interactions': avg_interactions,
+            'max_interactions': max_interactions,
+            'unique_tasks': unique_tasks
+        }
+
+
+def create_synthetic_dataset(num_trajectories: int = 10) -> TrajectoryDataset:
+    """
+    Create a synthetic dataset for testing purposes.
+    
+    Args:
+        num_trajectories: Number of trajectories to create
+        
+    Returns:
+        Synthetic trajectory dataset
+    """
+    dataset = TrajectoryDataset("synthetic_dataset")
+    
+    # Sample task descriptions
+    task_descriptions = [
+        "Book a flight from New York to London for next week",
+        "Find a vegetarian restaurant near downtown",
+        "Schedule a meeting with the marketing team for tomorrow",
+        "Order a new laptop with at least 16GB RAM",
+        "Write a congratulatory email to a colleague who got promoted",
+        "Research the best electric cars available in the market",
+        "Create a weekly meal plan with shopping list",
+        "Find information about tourist attractions in Barcelona",
+        "Help me debug a Python script that's giving an IndexError",
+        "Summarize the main points from the attached research paper"
+    ]
+    
+    # Create trajectories
+    for i in range(num_trajectories):
+        # Select task
+        task_idx = i % len(task_descriptions)
+        task = task_descriptions[task_idx]
+        
+        # Create interactions (2-4 turns)
+        num_turns = np.random.randint(2, 5)
+        interactions = []
+        
+        for j in range(num_turns):
+            if j == 0:
+                user_msg = f"I need help with this task: {task}"
+                agent_msg = f"I'd be happy to help you {task.lower()}. Could you provide more details about your preferences?"
+            elif j == num_turns - 1:
+                user_msg = "That sounds good. Please proceed with the final steps."
+                agent_msg = f"I've completed the task to {task.lower()}. Here's a summary of what I did..."
+            else:
+                user_msg = f"I prefer options that are {['affordable', 'convenient', 'high-quality'][j % 3]}."
+                agent_msg = f"Based on your preference for {['affordable', 'convenient', 'high-quality'][j % 3]} options, I recommend..."
+            
+            interactions.append({
+                'user': user_msg,
+                'agent': agent_msg
+            })
+        
+        # Determine if positive or negative example
+        is_positive = (i % 4 != 0)  # 75% positive, 25% negative
+        
+        # Create metadata
+        metadata = {
+            'is_positive': is_positive,
+            'quality_score': np.random.uniform(0.7, 0.9) if is_positive else np.random.uniform(0.3, 0.5),
+            'created_at': '2025-05-21'
+        }
+        
+        # Create and add trajectory
+        trajectory = Trajectory(
+            task_description=task,
+            interactions=interactions,
+            metadata=metadata
+        )
+        
+        dataset.add_trajectory(trajectory)
+    
+    return dataset