""" 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