swaleha19/agent_tuning_framework
0
1"""2Trajectory Data Management Module for Agent Tuning Optimization Framework3 4This module provides functionality for loading, processing, and managing agent interaction5trajectories for training and evaluation purposes.6"""7 8import os9import json10import pandas as pd11import numpy as np12from typing import List, Dict, Any, Union, Optional, Tuple13from tqdm import tqdm14 15class Trajectory:16 """Class representing a single agent interaction trajectory."""17 18 def __init__(19 self, 20 task_description: str,21 interactions: List[Dict[str, str]],22 metadata: Optional[Dict[str, Any]] = None23 ):24 """25 Initialize a trajectory.26 27 Args:28 task_description: Description of the task29 interactions: List of interaction turns (each with 'user' and 'agent' keys)30 metadata: Additional metadata about the trajectory31 """32 self.task_description = task_description33 self.interactions = interactions34 self.metadata = metadata or {}35 self.quality_score = self.metadata.get('quality_score', None)36 self.is_positive = self.metadata.get('is_positive', True)37 38 def to_dict(self) -> Dict[str, Any]:39 """40 Convert trajectory to dictionary.41 42 Returns:43 Dictionary representation of the trajectory44 """45 return {46 'task_description': self.task_description,47 'interactions': self.interactions,48 'metadata': self.metadata49 }50 51 @classmethod52 def from_dict(cls, data: Dict[str, Any]) -> 'Trajectory':53 """54 Create trajectory from dictionary.55 56 Args:57 data: Dictionary representation of the trajectory58 59 Returns:60 Trajectory instance61 """62 return cls(63 task_description=data['task_description'],64 interactions=data['interactions'],65 metadata=data.get('metadata', {})66 )67 68 def to_training_format(self, format_type: str = 'interleaved') -> str:69 """70 Convert trajectory to training format.71 72 Args:73 format_type: Format type ('interleaved', 'completion', etc.)74 75 Returns:76 Formatted trajectory as string77 """78 if format_type == 'interleaved':79 # Format as interleaved conversation80 result = f"Task: {self.task_description}\n\n"81 82 for i, interaction in enumerate(self.interactions):83 result += f"User: {interaction['user']}\n"84 result += f"Agent: {interaction['agent']}\n\n"85 86 return result.strip()87 88 elif format_type == 'completion':89 # Format as completion task (last agent response is the target)90 if not self.interactions:91 return ""92 93 result = f"Task: {self.task_description}\n\n"94 95 for i, interaction in enumerate(self.interactions[:-1]):96 result += f"User: {interaction['user']}\n"97 result += f"Agent: {interaction['agent']}\n\n"98 99 # Add last user query without agent response100 result += f"User: {self.interactions[-1]['user']}\n"101 result += f"Agent:"102 103 return result.strip(), self.interactions[-1]['agent'].strip()104 105 else:106 raise ValueError(f"Unsupported format type: {format_type}")107 108 def get_quality_score(self) -> float:109 """110 Get quality score for the trajectory.111 112 Returns:113 Quality score (0.0 to 1.0)114 """115 if self.quality_score is not None:116 return self.quality_score117 118 # Calculate simple quality score based on response length and complexity119 score = 0.0120 121 if not self.interactions:122 return score123 124 # Average response length (normalized)125 avg_length = np.mean([len(turn['agent']) for turn in self.interactions])126 length_score = min(avg_length / 500, 1.0) # Normalize to max of 500 chars127 128 # Response complexity (simple heuristic based on unique words)129 all_responses = " ".join([turn['agent'] for turn in self.interactions])130 unique_words = len(set(all_responses.lower().split()))131 complexity_score = min(unique_words / 200, 1.0) # Normalize to max of 200 unique words132 133 # Combine scores134 score = 0.6 * length_score + 0.4 * complexity_score135 136 # Cache the score137 self.quality_score = score138 self.metadata['quality_score'] = score139 140 return score141 142 143class TrajectoryDataset:144 """Dataset for managing collections of agent interaction trajectories."""145 146 def __init__(self, name: str):147 """148 Initialize the trajectory dataset.149 150 Args:151 name: Name of the dataset152 """153 self.name = name154 self.trajectories: List[Trajectory] = []155 self.positive_trajectories: List[Trajectory] = []156 self.negative_trajectories: List[Trajectory] = []157 158 def add_trajectory(self, trajectory: Trajectory) -> None:159 """160 Add a trajectory to the dataset.161 162 Args:163 trajectory: Trajectory to add164 """165 self.trajectories.append(trajectory)166 167 # Add to positive or negative list based on metadata168 if trajectory.is_positive:169 self.positive_trajectories.append(trajectory)170 else:171 self.negative_trajectories.append(trajectory)172 173 def load_from_json(self, file_path: str) -> None:174 """175 Load trajectories from JSON file.176 177 Args:178 file_path: Path to JSON file179 """180 with open(file_path, 'r') as f:181 data = json.load(f)182 183 if isinstance(data, list):184 # List of trajectories185 for item in data:186 self.add_trajectory(Trajectory.from_dict(item))187 elif isinstance(data, dict) and 'trajectories' in data:188 # Dictionary with trajectories key189 for item in data['trajectories']:190 self.add_trajectory(Trajectory.from_dict(item))191 else:192 raise ValueError(f"Unsupported JSON format in {file_path}")193 194 def save_to_json(self, file_path: str) -> None:195 """196 Save trajectories to JSON file.197 198 Args:199 file_path: Path to JSON file200 """201 data = {202 'name': self.name,203 'trajectories': [t.to_dict() for t in self.trajectories]204 }205 206 with open(file_path, 'w') as f:207 json.dump(data, f, indent=2)208 209 def get_trajectories(210 self, 211 positive_only: bool = False,212 negative_only: bool = False,213 min_quality: Optional[float] = None,214 max_samples: Optional[int] = None215 ) -> List[Trajectory]:216 """217 Get trajectories based on filtering criteria.218 219 Args:220 positive_only: Whether to return only positive trajectories221 negative_only: Whether to return only negative trajectories222 min_quality: Minimum quality score threshold223 max_samples: Maximum number of samples to return224 225 Returns:226 Filtered list of trajectories227 """228 if positive_only and negative_only:229 raise ValueError("Cannot set both positive_only and negative_only to True")230 231 # Select base list232 if positive_only:233 trajectories = self.positive_trajectories.copy()234 elif negative_only:235 trajectories = self.negative_trajectories.copy()236 else:237 trajectories = self.trajectories.copy()238 239 # Apply quality filter240 if min_quality is not None:241 trajectories = [t for t in trajectories if t.get_quality_score() >= min_quality]242 243 # Apply max samples limit244 if max_samples is not None and max_samples < len(trajectories):245 trajectories = trajectories[:max_samples]246 247 return trajectories248 249 def get_training_examples(250 self, 251 format_type: str = 'interleaved',252 positive_ratio: float = 0.8,253 min_quality: Optional[float] = 0.5,254 max_samples: Optional[int] = None255 ) -> Union[List[str], Tuple[List[str], List[str]]]:256 """257 Get formatted training examples from trajectories.258 259 Args:260 format_type: Format type ('interleaved', 'completion', etc.)261 positive_ratio: Ratio of positive to total examples262 min_quality: Minimum quality score threshold263 max_samples: Maximum number of samples to return264 265 Returns:266 Formatted training examples (format depends on format_type)267 """268 # Get positive and negative trajectories269 positive = self.get_trajectories(positive_only=True, min_quality=min_quality)270 negative = self.get_trajectories(negative_only=True)271 272 # Calculate sample counts273 if max_samples is not None:274 pos_count = int(max_samples * positive_ratio)275 neg_count = max_samples - pos_count276 else:277 pos_count = len(positive)278 neg_count = len(negative)279 280 # Sample trajectories281 if pos_count < len(positive):282 positive = np.random.choice(positive, pos_count, replace=False).tolist()283 284 if neg_count < len(negative):285 negative = np.random.choice(negative, neg_count, replace=False).tolist()286 287 # Format trajectories288 if format_type == 'interleaved':289 pos_examples = [t.to_training_format(format_type) for t in positive]290 neg_examples = [t.to_training_format(format_type) for t in negative]291 return pos_examples + neg_examples292 293 elif format_type == 'completion':294 pos_inputs = []295 pos_targets = []296 297 for t in positive:298 inp, target = t.to_training_format(format_type)299 pos_inputs.append(inp)300 pos_targets.append(target)301 302 neg_inputs = []303 neg_targets = []304 305 for t in negative:306 inp, target = t.to_training_format(format_type)307 neg_inputs.append(inp)308 neg_targets.append(target)309 310 return pos_inputs + neg_inputs, pos_targets + neg_targets311 312 else:313 raise ValueError(f"Unsupported format type: {format_type}")314 315 def analyze_dataset(self) -> Dict[str, Any]:316 """317 Analyze the dataset and return statistics.318 319 Returns:320 Dictionary of dataset statistics321 """322 if not self.trajectories:323 return {324 'total_trajectories': 0,325 'positive_count': 0,326 'negative_count': 0327 }328 329 # Basic counts330 total = len(self.trajectories)331 positive_count = len(self.positive_trajectories)332 negative_count = len(self.negative_trajectories)333 334 # Quality statistics335 quality_scores = [t.get_quality_score() for t in self.trajectories]336 avg_quality = np.mean(quality_scores)337 min_quality = np.min(quality_scores)338 max_quality = np.max(quality_scores)339 340 # Interaction statistics341 interaction_counts = [len(t.interactions) for t in self.trajectories]342 avg_interactions = np.mean(interaction_counts)343 max_interactions = np.max(interaction_counts)344 345 # Task diversity (simple heuristic based on unique task descriptions)346 unique_tasks = len(set([t.task_description for t in self.trajectories]))347 348 return {349 'total_trajectories': total,350 'positive_count': positive_count,351 'negative_count': negative_count,352 'positive_ratio': positive_count / total if total > 0 else 0,353 'avg_quality': avg_quality,354 'min_quality': min_quality,355 'max_quality': max_quality,356 'avg_interactions': avg_interactions,357 'max_interactions': max_interactions,358 'unique_tasks': unique_tasks359 }360 361 362def create_synthetic_dataset(num_trajectories: int = 10) -> TrajectoryDataset:363 """364 Create a synthetic dataset for testing purposes.365 366 Args:367 num_trajectories: Number of trajectories to create368 369 Returns:370 Synthetic trajectory dataset371 """372 dataset = TrajectoryDataset("synthetic_dataset")373 374 # Sample task descriptions375 task_descriptions = [376 "Book a flight from New York to London for next week",377 "Find a vegetarian restaurant near downtown",378 "Schedule a meeting with the marketing team for tomorrow",379 "Order a new laptop with at least 16GB RAM",380 "Write a congratulatory email to a colleague who got promoted",381 "Research the best electric cars available in the market",382 "Create a weekly meal plan with shopping list",383 "Find information about tourist attractions in Barcelona",384 "Help me debug a Python script that's giving an IndexError",385 "Summarize the main points from the attached research paper"386 ]387 388 # Create trajectories389 for i in range(num_trajectories):390 # Select task391 task_idx = i % len(task_descriptions)392 task = task_descriptions[task_idx]393 394 # Create interactions (2-4 turns)395 num_turns = np.random.randint(2, 5)396 interactions = []397 398 for j in range(num_turns):399 if j == 0:400 user_msg = f"I need help with this task: {task}"401 agent_msg = f"I'd be happy to help you {task.lower()}. Could you provide more details about your preferences?"402 elif j == num_turns - 1:403 user_msg = "That sounds good. Please proceed with the final steps."404 agent_msg = f"I've completed the task to {task.lower()}. Here's a summary of what I did..."405 else:406 user_msg = f"I prefer options that are {['affordable', 'convenient', 'high-quality'][j % 3]}."407 agent_msg = f"Based on your preference for {['affordable', 'convenient', 'high-quality'][j % 3]} options, I recommend..."408 409 interactions.append({410 'user': user_msg,411 'agent': agent_msg412 })413 414 # Determine if positive or negative example415 is_positive = (i % 4 != 0) # 75% positive, 25% negative416 417 # Create metadata418 metadata = {419 'is_positive': is_positive,420 'quality_score': np.random.uniform(0.7, 0.9) if is_positive else np.random.uniform(0.3, 0.5),421 'created_at': '2025-05-21'422 }423 424 # Create and add trajectory425 trajectory = Trajectory(426 task_description=task,427 interactions=interactions,428 metadata=metadata429 )430 431 dataset.add_trajectory(trajectory)432 433 return dataset434