#!/usr/bin/env python3 """ Main experiment orchestrator for personalization benchmark. This script runs all baselines and the proposed methods with PROPER multi-turn conversation simulation, user preference enforcement, and LLM-based evaluation. Usage: python run_experiments.py --config config.yaml python run_experiments.py --methods vanilla,rag,rag_vector --datasets gpqa,aime """ import argparse import json import yaml import os import sys import numpy as np from pathlib import Path from datetime import datetime from typing import List, Dict, Any, Optional from dataclasses import dataclass, asdict import logging import re from concurrent.futures import ThreadPoolExecutor, as_completed import threading import time # Add paths sys.path.insert(0, str(Path(__file__).parent.parent)) sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from datasets_extended import get_dataset, get_all_datasets, get_challenging_datasets from evaluation.llm_judge import LLMJudge, BatchEvaluator, ConversationMetrics from conflict_scenario_generator import ConflictScenarioGenerator from adapters.personalized_llm_adapter import PersonalizedLLMAdapter, create_baseline_adapter from agents.local_user_agent import LocalUserAgent, SharedLocalUserAgent, TERMINATION_SIGNAL from agents.vllm_user_agent import VLLMUserAgent, VLLMAgentClient from agents.openai_user_agent import OpenAIUserAgent from agents.batch_vllm_agent import BatchConversationGenerator, BatchVLLMClient logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) @dataclass class ExperimentConfig: """Configuration for an experiment run.""" # Methods to compare methods: List[str] # Datasets to use datasets: List[str] # User profiles n_profiles: int = 200 profile_path: Optional[str] = None # Profile range (for splitting jobs) start_profile: int = 0 # Inclusive, 0-indexed end_profile: Optional[int] = None # Exclusive, None means all # Session settings n_sessions_per_profile: int = 30 max_turns_per_session: int = 15 # Increased for harder tasks # Model settings user_model: str = "meta-llama/Llama-3.3-70B-Instruct" agent_model: str = "meta-llama/Llama-3.1-8B-Instruct" judge_model: str = "meta-llama/Llama-3.3-70B-Instruct" # Output settings output_dir: str = "results" save_conversations: bool = True # Conflict testing conflict_ratio: float = 0.3 # proportion of queries that trigger conflicts # Compute settings batch_size: int = 4 n_gpus: int = 4 # vLLM settings (for high-performance inference) use_vllm: bool = False vllm_user_url: str = "http://localhost:8004/v1" # 70B user simulator vllm_agent_url: str = "http://localhost:8003/v1" # 8B agent # OpenAI user simulator (alternative to vLLM user agent) use_openai_user: bool = False openai_user_model: str = "gpt-5" # Model name for OpenAI user agent # Reward mode: "keyword" (implicit user signals), "llm" (GPT-4o-mini), or "llm_local" (local vLLM) # This is a global option applied to ALL methods that use RL updates reward_mode: str = "keyword" # vLLM URL for local reward model (only used when reward_mode="llm_local") reward_vllm_url: str = "http://localhost:8005/v1" # Parallel/Batch processing parallel_profiles: int = 50 # Number of profiles to process in parallel use_batch_processing: bool = True # Use turn-synchronous batch processing for vanilla/all_memory batch_size_conversations: int = 50 # Number of conversations to batch together # Continue from existing experiment (for extending sessions) continue_from: Optional[str] = None # Path to existing output directory to continue from # Available methods AVAILABLE_METHODS = { "vanilla": "No memory, no personalization", "contextual": "Full history in context, summarize when overflow", "reflection": "CollaborativeAgents' agent_notes approach", "reflection_grpo": "Reflection + GRPO training", "all_memory": "All extracted memories in context (no retrieval)", "rag": "Extractor + RAG (no user vector)", "rag_dynamic": "Extractor + RAG with dynamic topk (min=3, max=8, ratio=0.5)", "rag_rewrite": "Extractor + RAG with LLM preference rewrite/merge", "rag_rewrite_vector": "Extractor + RAG + user vector + LLM preference rewrite", "rag_vector": "Extractor + RAG + user vector (proposed method)", "rag_vector_no_short": "Ablation: RAG + user vector without z_short (long-term only)", "rag_vector_no_long": "Ablation: RAG + user vector without z_long (short-term only)", "rag_vector_fast": "Extractor + RAG + user vector with 10x learning rate", "rag_vector_consolidate": "Extractor + RAG + user vector with session-level preference consolidation", "rag_vector_balanced": "Extractor + RAG + user vector with balanced rewards (10x LR + positive signal for good turns)", "rag_bge": "Extractor + RAG with BGE reranker (278M)", "rag_vector_bge": "Extractor + RAG + user vector with BGE reranker (278M)", } class ExperimentRunner: """Main experiment runner.""" def __init__(self, config: ExperimentConfig): self.config = config # Use existing directory if continuing, otherwise create new timestamped one if config.continue_from: self.output_dir = Path(config.continue_from) if not self.output_dir.exists(): raise ValueError(f"Continue-from directory does not exist: {config.continue_from}") logger.info(f"Continuing from existing experiment: {self.output_dir}") else: self.output_dir = Path(config.output_dir) / datetime.now().strftime("%Y%m%d_%H%M%S") self.output_dir.mkdir(parents=True, exist_ok=True) # Save/update config with open(self.output_dir / "config.yaml", "w") as f: yaml.dump(asdict(config), f) # Initialize components self.judge = LLMJudge(model_name=config.judge_model) self.batch_evaluator = BatchEvaluator(self.judge) self.conflict_generator = ConflictScenarioGenerator() # Load datasets self.datasets = {} for ds_name in config.datasets: try: self.datasets[ds_name] = get_dataset(ds_name) logger.info(f"Loaded dataset: {ds_name}") except Exception as e: logger.warning(f"Failed to load dataset {ds_name}: {e}") # Load or generate profiles self.profiles = self._load_profiles() def _load_profiles(self) -> List[Dict]: """Load user profiles from file or generate.""" logger.info(f"Profile path configured: {self.config.profile_path}") if self.config.profile_path: profile_path = Path(self.config.profile_path) if profile_path.exists(): profiles = [] with open(profile_path) as f: for line in f: line = line.strip() if line: profiles.append(json.loads(line)) logger.info(f"Loaded {len(profiles)} profiles from {self.config.profile_path}") return profiles[:self.config.n_profiles] else: logger.warning(f"Profile path does not exist: {self.config.profile_path}") # Generate simple placeholder profiles if no file provided logger.info(f"Generating {self.config.n_profiles} placeholder profiles...") profiles = [] for i in range(self.config.n_profiles): profiles.append({ "id": i, "persona": f"User {i+1} is a curious individual seeking help with problem solving.", "preferences": [ "Provide clear, step-by-step explanations", "Use simple language when possible", "Give examples to illustrate concepts", "Be concise but thorough", "Acknowledge when something is uncertain" ] }) # Save generated profiles profile_path = self.output_dir / "generated_profiles.json" with open(profile_path, "w") as f: json.dump(profiles, f, indent=2) logger.info(f"Generated and saved {len(profiles)} placeholder profiles") return profiles def _create_method_adapter(self, method: str, profile: Dict, use_shared_models: bool = False) -> Any: """Create adapter for a specific method. Args: method: One of the baseline method names profile: User profile dict (used later in start_session, not constructor) use_shared_models: If True, share embedding/reranker models across parallel workers. ESSENTIAL for parallel profile processing to avoid OOM. Returns: Configured adapter instance """ # Auto-detect available GPUs and set device assignment accordingly # Layout with local 70B user (4 GPUs): # GPU 0-1: 70B user simulator (TP=2) # GPU 2: 8B agent vLLM server # GPU 3: Embedding + Reranker + Extractor # Layout with OpenAI user (2 GPUs): # GPU 0: 8B agent vLLM server # GPU 1: Embedding + Reranker + Extractor device_assignment = None try: import torch n_gpus = torch.cuda.device_count() if n_gpus >= 4: # 4 GPU layout: 70B user on 0-1, agent on 2, adapters on 3 device_assignment = { "embed": "cuda:3", "reranker": "cuda:3", "extractor": "cuda:3", } elif n_gpus >= 2: # 2 GPU layout: agent on 0, adapters on 1 device_assignment = { "embed": "cuda:1", "reranker": "cuda:1", "extractor": "cuda:1", } elif n_gpus == 1: device_assignment = { "embed": "cuda:0", "reranker": "cuda:0", "extractor": "cuda:0", } except ImportError: pass adapter = create_baseline_adapter( method, device_assignment=device_assignment, use_vllm=self.config.use_vllm, use_shared_models=use_shared_models, reward_mode=self.config.reward_mode, reward_vllm_url=self.config.reward_vllm_url, ) # Profile will be passed to start_session() when the conversation begins return adapter def _export_user_vectors(self, method: str, adapters: Dict[int, Any]) -> None: """ Export user vectors from all adapters to disk for later analysis. Saves both .npz (efficient numpy format) and .json (human-readable). Args: method: Method name for the output directory adapters: Dict mapping profile_idx to adapter instances """ method_dir = self.output_dir / method # Collect all user vectors from adapters all_vectors = {} for profile_idx, adapter in adapters.items(): if hasattr(adapter, 'export_all_user_vectors'): vectors = adapter.export_all_user_vectors() all_vectors.update(vectors) if not all_vectors: logger.info(f" No user vectors to export for {method}") return # Save as .npz for efficient analysis npz_path = method_dir / "user_vectors.npz" user_ids = list(all_vectors.keys()) k = len(all_vectors[user_ids[0]]["z_long"]) z_long = np.zeros((len(user_ids), k), dtype=np.float32) z_short = np.zeros((len(user_ids), k), dtype=np.float32) reward_ma = np.zeros(len(user_ids), dtype=np.float32) for i, uid in enumerate(user_ids): z_long[i] = all_vectors[uid]["z_long"] z_short[i] = all_vectors[uid]["z_short"] reward_ma[i] = all_vectors[uid]["reward_ma"] np.savez( npz_path, user_ids=np.array(user_ids), z_long=z_long, z_short=z_short, reward_ma=reward_ma, ) # Also save summary stats as JSON summary = { "n_users": len(user_ids), "vector_dim": k, "z_long_norms": {uid: all_vectors[uid]["z_long_norm"] for uid in user_ids}, "z_short_norms": {uid: all_vectors[uid]["z_short_norm"] for uid in user_ids}, "reward_mas": {uid: all_vectors[uid]["reward_ma"] for uid in user_ids}, "stats": { "z_long_norm_mean": float(np.mean([all_vectors[uid]["z_long_norm"] for uid in user_ids])), "z_long_norm_max": float(np.max([all_vectors[uid]["z_long_norm"] for uid in user_ids])), "z_long_norm_std": float(np.std([all_vectors[uid]["z_long_norm"] for uid in user_ids])), } } with open(method_dir / "user_vectors_summary.json", "w") as f: json.dump(summary, f, indent=2) logger.info(f" Exported {len(user_ids)} user vectors to {npz_path}") def run_single_session( self, method: str, profile: Dict, problem: Dict, is_conflict_query: bool = False, adapter: Any = None, user_agent: Any = None ) -> Dict: """Run a single session with PROPER multi-turn conversation and user simulation. This implements: 1. User simulator that role-plays with preferences 2. Multi-turn conversation (up to max_turns) 3. Preference enforcement by simulated user 4. Proper metrics extraction from conversation """ # Use provided adapter (reused across sessions) or create new one agent_adapter = adapter if adapter else self._create_method_adapter(method, profile) # Prepare conflict scenario if needed conflict_scenario = None original_problem = problem.get("problem", problem.get("question", "")) if is_conflict_query: conflict_scenario = self.conflict_generator.generate_for_profile( profile["preferences"], problem.get("domain", "general") ) if conflict_scenario: problem = dict(problem) problem["problem"] = conflict_scenario["query"] query = problem.get("problem", problem.get("question", "")) # Extract user preferences as formatted string user_prefs = profile.get("preferences", []) if isinstance(user_prefs, list) and len(user_prefs) > 0: if isinstance(user_prefs[0], dict): # Structured preferences with condition/action pref_str = "\n".join([ f"- When {p.get('condition', '')}, {p.get('action', '')}" for p in user_prefs ]) else: # Simple string preferences pref_str = "\n".join([f"- {p}" for p in user_prefs]) else: pref_str = str(user_prefs) user_persona = profile.get("persona", "A user seeking help with problem solving.") # Create user agent for this session (or reuse provided one) if user_agent is None: if self.config.use_openai_user: user_agent = OpenAIUserAgent( user_task_description="Help the user solve their problem.", problem=query, user_persona=user_persona, user_preferences=pref_str, model=self.config.openai_user_model, ) elif self.config.use_vllm: user_agent = VLLMUserAgent( user_task_description="Help the user solve their problem.", problem=query, user_persona=user_persona, user_preferences=pref_str, vllm_url=self.config.vllm_user_url, ) else: user_agent = SharedLocalUserAgent( user_task_description="Help the user solve their problem.", problem=query, user_persona=user_persona, user_preferences=pref_str, ) # Initialize conversation turns = [] full_user_log = [] # Detailed user agent outputs # Metrics tracking enforcement_count = 0 disappointment_count = 0 user_token_count = 0 agent_token_count = 0 preference_compliance_scores = [] try: # Initialize adapter for this user if hasattr(agent_adapter, 'initialize'): agent_adapter.initialize() if hasattr(agent_adapter, 'start_session'): agent_adapter.start_session( user_id=profile.get("user_id", "test_user"), user_profile={"preferences": user_prefs, "persona": user_persona} ) # Start with agent greeting conversation = [{"role": "assistant", "content": "How can I help you today?"}] # Multi-turn conversation loop for turn_num in range(self.config.max_turns_per_session): # === User Turn === user_response = user_agent.generate_user_response(conversation) if user_response is None: logger.warning(f"User agent failed to respond at turn {turn_num}") break user_message = str(user_response.get("response", "")) user_token_count += len(user_message.split()) # Add to conversation conversation.append({"role": "user", "content": user_message}) turns.append({"role": "user", "content": user_message}) full_user_log.append(user_response) # Check for termination if user_response.get("should_terminate", False) or TERMINATION_SIGNAL in user_message: break # Detect preference enforcement (user correcting agent) enforcement_keywords = ["please", "i asked", "i said", "i prefer", "can you", "could you", "instead"] if any(kw in user_message.lower() for kw in enforcement_keywords): enforcement_count += 1 # === Agent Turn === if hasattr(agent_adapter, 'generate_response'): response = agent_adapter.generate_response(user_message, conversation[:-1]) agent_content = response.get("response", str(response)) if isinstance(response, dict) else str(response) elif callable(agent_adapter): agent_content = agent_adapter(conversation) else: agent_content = "[Error: Adapter not properly configured]" agent_token_count += len(agent_content.split()) # Add to conversation conversation.append({"role": "assistant", "content": agent_content}) turns.append({"role": "assistant", "content": agent_content}) # Estimate preference compliance for this turn (heuristic based on user satisfaction) # If user doesn't enforce in next turn, assume compliance # This is a simplified heuristic - LLM judge would be more accurate compliance_score = 0.8 if enforcement_count == 0 else max(0.2, 1.0 - 0.2 * enforcement_count) preference_compliance_scores.append(compliance_score) # End session if hasattr(agent_adapter, 'end_session'): adapter_metrics = agent_adapter.end_session(task_success=True) else: adapter_metrics = {} except Exception as e: import traceback logger.error(f"Error in session: {e}") logger.error(f"Full traceback:\n{traceback.format_exc()}") turns.append({"role": "assistant", "content": f"[Error: {e}]"}) # Compute metrics total_turns = len(turns) total_token_count = user_token_count + agent_token_count # Check if user reached a satisfactory answer (from last user response) task_success = False if full_user_log: last_user = full_user_log[-1] if last_user.get("should_terminate", False): draft = last_user.get("draft_answer", "") # Consider success if draft answer is not empty/"I don't know" task_success = bool(draft) and draft.lower() != "i don't know" # Compute average compliance avg_compliance = sum(preference_compliance_scores) / len(preference_compliance_scores) if preference_compliance_scores else 0.5 # Conflict resolution (if this was a conflict test) conflict_accuracy = 0.0 if is_conflict_query and conflict_scenario: # Check if the correct preference was applied expected_pref = conflict_scenario.get("expected_preference", "") # Simple heuristic: check if expected preference keywords appear in agent responses agent_texts = " ".join([t["content"] for t in turns if t["role"] == "assistant"]) if expected_pref and any(kw in agent_texts.lower() for kw in expected_pref.lower().split()[:3]): conflict_accuracy = 1.0 # Over-personalization detection (heuristic: if agent mentions preferences not in profile) over_personalization = 0.0 metrics = ConversationMetrics( task_success=task_success, turns_to_success=total_turns if task_success else -1, total_turns=total_turns, user_token_count=user_token_count, enforcement_count=enforcement_count, disappointment_count=disappointment_count, total_token_count=total_token_count, agent_token_count=agent_token_count, preference_compliance_scores=preference_compliance_scores, conflict_resolution_accuracy=conflict_accuracy, over_personalization_rate=over_personalization, ) return { "method": method, "profile_id": profile.get("user_id", "unknown"), "problem_id": problem.get("problem_id", str(hash(query))[:8]), "problem": original_problem, "ground_truth_solution": problem.get("solution", problem.get("answer", "")), "is_conflict_test": is_conflict_query, "conflict_scenario": conflict_scenario, "conversation": {"turns": turns} if self.config.save_conversations else None, "full_user_log": full_user_log if self.config.save_conversations else None, "metrics": asdict(metrics), "adapter_metrics": adapter_metrics if 'adapter_metrics' in dir() else {}, } def _run_profile_sessions( self, method: str, profile_idx: int, profile: Dict, adapter: Any = None ) -> List[Dict]: """Run all sessions for a single profile. Thread-safe for parallel execution.""" profile_results = [] # Create vLLM-based agent client if using vLLM (for methods that need it) vllm_agent = None if self.config.use_vllm and method == "vanilla": vllm_agent = VLLMAgentClient( vllm_url=self.config.vllm_agent_url, system_prompt="You are a helpful AI assistant for problem-solving tasks." ) # Run sessions across datasets session_idx = 0 for ds_name, dataset in self.datasets.items(): samples = dataset.get_testset() for sample in samples: if session_idx >= self.config.n_sessions_per_profile: break # Decide if this is a conflict query is_conflict = (session_idx % int(1 / self.config.conflict_ratio)) == 0 problem = { "problem": sample.problem, "solution": sample.solution, "problem_id": sample.problem_id, "domain": sample.domain, } try: result = self.run_single_session( method=method, profile=profile, problem=problem, is_conflict_query=is_conflict, adapter=vllm_agent if vllm_agent else adapter ) profile_results.append(result) except Exception as e: logger.error(f"Error in session for profile {profile_idx}: {e}") session_idx += 1 return profile_results def run_method(self, method: str) -> List[Dict]: """Run all sessions for a single method with checkpointing and parallel processing.""" logger.info(f"Running method: {method}") # Setup method directory and checkpoint method_dir = self.output_dir / method method_dir.mkdir(exist_ok=True) checkpoint_file = method_dir / "checkpoint.json" results_file = method_dir / "results.json" # Load existing results and checkpoint results = [] completed_profiles = set() sessions_per_profile = {} # Track session count per profile for continue functionality if checkpoint_file.exists(): with open(checkpoint_file, "r") as f: checkpoint = json.load(f) completed_profiles = set(checkpoint.get("completed_profiles", [])) sessions_per_profile = checkpoint.get("sessions_per_profile", {}) logger.info(f" Resuming from checkpoint: {len(completed_profiles)} profiles completed") if sessions_per_profile: total_sessions = sum(sessions_per_profile.values()) logger.info(f" Session-level tracking: {total_sessions} sessions across {len(sessions_per_profile)} profiles") if results_file.exists(): with open(results_file, "r") as f: results = json.load(f) # Determine profile range start_idx = self.config.start_profile end_idx = self.config.end_profile if self.config.end_profile else len(self.profiles) # Build list of profiles that need more sessions profiles_to_run = [] for idx in range(start_idx, min(end_idx, len(self.profiles))): existing_sessions = sessions_per_profile.get(str(idx), 0) if existing_sessions < self.config.n_sessions_per_profile: profiles_to_run.append(idx) # Log what we're running if sessions_per_profile: total_existing = sum(sessions_per_profile.get(str(idx), 0) for idx in profiles_to_run) total_needed = len(profiles_to_run) * self.config.n_sessions_per_profile logger.info(f" Running profiles {start_idx} to {end_idx-1}: {len(profiles_to_run)} profiles need sessions") logger.info(f" Sessions: {total_existing} existing, {total_needed - total_existing} remaining") else: logger.info(f" Running profiles {start_idx} to {end_idx-1} ({len(profiles_to_run)} remaining)") # When using batch processing with vLLM or OpenAI user: use turn-synchronous batch mode # This batches both user and agent calls for maximum throughput if self.config.use_batch_processing and self.config.use_vllm: user_type = "OpenAI" if self.config.use_openai_user else "local vLLM" logger.info(f" Using BATCH processing ({user_type} user) for {method}") return self._run_method_batch( method, profiles_to_run, results, completed_profiles, sessions_per_profile, checkpoint_file, results_file ) # Decide on parallelization for sequential methods n_parallel = self.config.parallel_profiles if (self.config.use_vllm or self.config.use_openai_user) else 1 if n_parallel > 1: logger.info(f" Using parallel processing with {n_parallel} workers") self._run_method_parallel( method, profiles_to_run, results, completed_profiles, sessions_per_profile, checkpoint_file, results_file ) else: # Sequential execution (original behavior) # Create ONE adapter per method and reuse it (avoids GPU OOM from repeated model loading) adapter = self._create_method_adapter(method, None) adapter.initialize() for profile_idx in profiles_to_run: profile = self.profiles[profile_idx] logger.info(f" Profile {profile_idx + 1}/{len(self.profiles)}") profile_results = self._run_profile_sessions(method, profile_idx, profile, adapter) # Add profile results to overall results results.extend(profile_results) completed_profiles.add(profile_idx) sessions_per_profile[str(profile_idx)] = self.config.n_sessions_per_profile # Save checkpoint and results after each profile with open(checkpoint_file, "w") as f: json.dump({ "completed_profiles": sorted(list(completed_profiles)), "sessions_per_profile": sessions_per_profile }, f) with open(results_file, "w") as f: json.dump(results, f, indent=2) logger.info(f" Profile {profile_idx + 1} completed and checkpointed") # Export user vectors at the end of sequential processing self._export_user_vectors(method, {0: adapter}) return results def _run_method_parallel( self, method: str, profiles_to_run: List[int], results: List[Dict], completed_profiles: set, sessions_per_profile: Dict[str, int], checkpoint_file: Path, results_file: Path ): """Run profiles in parallel using ThreadPoolExecutor. Uses shared model singletons for embedding/reranker to avoid OOM when multiple workers try to load their own copies. """ n_parallel = self.config.parallel_profiles results_lock = threading.Lock() start_time = time.time() profiles_completed = 0 def process_profile(profile_idx: int) -> tuple: """Process a single profile and return (profile_idx, results).""" profile = self.profiles[profile_idx] # Create adapter with shared models to avoid OOM from duplicate model loading adapter = self._create_method_adapter(method, profile, use_shared_models=True) profile_results = self._run_profile_sessions(method, profile_idx, profile, adapter) return profile_idx, profile_results with ThreadPoolExecutor(max_workers=n_parallel) as executor: # Submit all profile jobs future_to_profile = { executor.submit(process_profile, idx): idx for idx in profiles_to_run } # Process completed profiles for future in as_completed(future_to_profile): profile_idx = future_to_profile[future] try: idx, profile_results = future.result() with results_lock: results.extend(profile_results) completed_profiles.add(idx) sessions_per_profile[str(idx)] = self.config.n_sessions_per_profile profiles_completed += 1 # Save checkpoint with session-level tracking with open(checkpoint_file, "w") as f: json.dump({ "completed_profiles": sorted(list(completed_profiles)), "sessions_per_profile": sessions_per_profile }, f) with open(results_file, "w") as f: json.dump(results, f, indent=2) # Log progress with throughput estimate elapsed = time.time() - start_time profiles_per_hour = profiles_completed / elapsed * 3600 if elapsed > 0 else 0 sessions_per_hour = len(results) / elapsed * 3600 if elapsed > 0 else 0 logger.info( f" Profile {idx + 1} completed " f"({profiles_completed}/{len(profiles_to_run)}) - " f"{profiles_per_hour:.1f} profiles/hr, {sessions_per_hour:.1f} sessions/hr" ) except Exception as e: logger.error(f" Profile {profile_idx} failed: {e}") # Note: Parallel mode doesn't export user vectors because adapters are # created/destroyed per profile. Use batch mode for vector export. logger.info(f" Parallel mode: user vectors not exported (use batch mode)") def _run_method_batch( self, method: str, profiles_to_run: List[int], results: List[Dict], completed_profiles: set, sessions_per_profile: Dict[str, int], checkpoint_file: Path, results_file: Path ) -> List[Dict]: """ Turn-synchronous batch processing for ALL methods. At each turn, user calls are batched concurrently via AsyncOpenAI, then agent responses go through personalization adapters. Sessions within a profile run sequentially (for stateful memory). """ from agents.batch_vllm_agent import BatchOpenAIClient, BatchVLLMClient, TERMINATION_SIGNAL from json_repair import repair_json start_time = time.time() # Create user client (OpenAI API or local vLLM) if self.config.use_openai_user: user_client = BatchOpenAIClient( model=self.config.openai_user_model, max_tokens=4096, max_concurrent=32, api_key=os.environ.get("OPENAI_API_KEY"), ) logger.info(f" Using OpenAI user simulator: {self.config.openai_user_model}") else: user_client = BatchVLLMClient( vllm_url=self.config.vllm_user_url, max_tokens=1024, # User responses typically short, but allow for edge cases temperature=1.0, timeout=None, max_concurrent=100, json_mode=True, # User simulator needs JSON output ) logger.info(f" Using local vLLM user simulator: {self.config.vllm_user_url}") # Create async agent client for batched vLLM calls agent_client = BatchVLLMClient( vllm_url=self.config.vllm_agent_url, max_tokens=2048, temperature=0.7, timeout=None, # Infinite timeout for long generations max_concurrent=100, ) USER_PROMPT_TEMPLATE = ( "You are a user simulator collaborating with an agent to solve a problem. " "You will be provided with a problem description, and you must get the agent to help you solve it. " "You will also be provided with user preferences, which you must follow and actively enforce throughout the conversation.\n\n" "# Problem Description\n{problem}\nNote: the agent cannot see this problem description.\n\n" "# User Persona\n{user_persona}\n\n" "# User Preferences\n{user_preferences}\n" "These preferences are NON-NEGOTIABLE that define how you prefer the agent to behave. They must be strictly enforced:\n" " - **Answer clarifying questions**: The agent may ask clarifying questions before attempting an answer. " "Answer such questions, and do not enforce preferences about answer format or content while the agent is clarifying.\n" " - **Enforce immediately**: Every agent response must satisfy your preferences before you can proceed. " "Explicitly ask the agent to adjust their response until it complies.\n" " - **Never proceed without compliance**: Do NOT update your draft answer, do NOT consider terminating, " "and do NOT move forward until the agent follows your preferences.\n\n" "# Draft Answer Management\n" "- **Maintain a working draft**: Start with \"I don't know\". Update your draft answer based on what you learn from agent responses.\n" "- **Don't update when enforcing preferences**: If the agent response does not follow your preferences, " "do NOT update your draft answer, regardless of whether the agent provides helpful information.\n\n" "# Conversation Termination\n" "Before generating your response, determine if you should terminate:\n" " - Do you feel like your draft answer is a good answer to the problem?\n" " - Do you feel like the agent cannot help further?\n" "If the agent response does not follow your preferences, you must NOT terminate - instead, enforce the preferences.\n" "When ready to terminate, respond with \"TERMINATE\".\n\n" "# Output Format (respond in JSON):\n" "{{\n" " \"preferences_check\": \"For EACH relevant preference, evaluate: is it satisfied?\",\n" " \"enforce_preferences\": true/false,\n" " \"reasoning\": \"Brief reasoning (2-3 sentences). Does agent follow preferences? If no, enforce. If yes, update draft.\",\n" " \"draft_answer\": \"Your current working draft answer\",\n" " \"should_terminate\": true/false,\n" " \"response\": \"Your response to the agent\"\n" "}}" ) def parse_user_response(content): if not content: return None try: parsed = repair_json(content, return_objects=True) if isinstance(parsed, dict) and "response" in parsed: return parsed except: pass if TERMINATION_SIGNAL in (content or ""): return {"reasoning": "", "draft_answer": "", "should_terminate": True, "response": TERMINATION_SIGNAL} return {"reasoning": "", "draft_answer": "", "should_terminate": False, "response": content or ""} def reverse_roles(conversation): return [ {"role": "user" if m["role"] == "assistant" else "assistant", "content": m["content"]} for m in conversation ] # Create per-profile adapters adapters = {} profile_sessions = {} # Build session problem list ONCE (shared across all profiles for controlled comparison) # Each dataset contributes exactly n_per_dataset problems (front 10), no repeats shared_sessions = [] dataset_names = list(self.datasets.keys()) n_per_dataset = self.config.n_sessions_per_profile // len(dataset_names) remainder = self.config.n_sessions_per_profile % len(dataset_names) for i, ds_name in enumerate(dataset_names): ds_obj = self.datasets[ds_name] items = ds_obj.get_testset() n_take = n_per_dataset + (1 if i < remainder else 0) if n_take > len(items): logger.warning(f" Dataset {ds_name} has only {len(items)} problems, need {n_take}") for j in range(n_take): item = items[j % len(items)] shared_sessions.append({"problem": item.problem, "solution": item.solution, "domain": ds_obj.domain}) n_conflict = int(len(shared_sessions) * self.config.conflict_ratio) shared_session_list = [(s, idx < n_conflict) for idx, s in enumerate(shared_sessions)] logger.info(f" Built shared session list: {len(shared_sessions)} problems from {len(dataset_names)} datasets ({n_per_dataset} each, same for all profiles)") for profile_idx in profiles_to_run: profile = self.profiles[profile_idx] adapter = self._create_method_adapter(method, profile, use_shared_models=True) if hasattr(adapter, 'initialize'): adapter.initialize() adapters[profile_idx] = adapter profile_sessions[profile_idx] = shared_session_list n_sessions = self.config.n_sessions_per_profile # Calculate sessions to run per profile (accounting for existing sessions) sessions_to_run_per_profile = {} for profile_idx in profiles_to_run: existing = sessions_per_profile.get(str(profile_idx), 0) remaining = n_sessions - existing if remaining > 0: sessions_to_run_per_profile[profile_idx] = (existing, remaining) # (start_session, count) if sessions_to_run_per_profile: total_remaining = sum(v[1] for v in sessions_to_run_per_profile.values()) logger.info(f" Batch: {len(sessions_to_run_per_profile)} profiles, {total_remaining} sessions remaining") else: logger.info(f" Batch: All sessions already completed") return results # Process sessions in rounds for session_idx in range(n_sessions): # Initialize all conversations for this round all_states = {} # profile_idx -> state dict active_set = set() for profile_idx in profiles_to_run: # Skip if this profile doesn't need this session if profile_idx not in sessions_to_run_per_profile: continue start_session, _ = sessions_to_run_per_profile[profile_idx] if session_idx < start_session: continue # Already completed this session if session_idx >= len(profile_sessions[profile_idx]): continue problem_dict, is_conflict = profile_sessions[profile_idx][session_idx] profile = self.profiles[profile_idx] query = problem_dict["problem"] if is_conflict: cs = self.conflict_generator.generate_for_profile( profile.get("preferences", []), problem_dict.get("domain", "general")) if cs: query = cs["query"] user_prefs = profile.get("preferences", []) if isinstance(user_prefs, list) and user_prefs: if isinstance(user_prefs[0], dict): pref_str = "\n".join([f"- When {p.get('condition','')}, {p.get('action','')}" for p in user_prefs]) else: pref_str = "\n".join([f"- {p}" for p in user_prefs]) else: pref_str = str(user_prefs) user_persona = profile.get("persona", "A user seeking help with problem solving.") adapter = adapters[profile_idx] if hasattr(adapter, 'start_session'): adapter.start_session( user_id=profile.get("user_id", f"user_{profile_idx}"), user_profile={"preferences": user_prefs, "persona": user_persona} ) all_states[profile_idx] = { "conversation": [{"role": "assistant", "content": "How can I help you today?"}], "full_log": [], "system_prompt": USER_PROMPT_TEMPLATE.format( problem=query, user_persona=user_persona, user_preferences=pref_str), "problem_dict": problem_dict, "is_conflict": is_conflict, "enforcement_count": 0, } active_set.add(profile_idx) # Turn-synchronous loop for turn in range(self.config.max_turns_per_session): if not active_set: break # Batch user calls active_list = sorted(active_set) user_msgs_batch = [] for pidx in active_list: state = all_states[pidx] msgs = [{"role": "system", "content": state["system_prompt"]}] msgs.extend(reverse_roles(state["conversation"])) user_msgs_batch.append(msgs) user_responses = user_client.batch_completion(user_msgs_batch) # Process user responses and prepare agent prompts for batching to_remove = [] agent_prompts_batch = [] # List of (pidx, messages, context) for i, pidx in enumerate(active_list): state = all_states[pidx] parsed = parse_user_response(user_responses[i]) if parsed is None: to_remove.append(pidx) continue user_msg = str(parsed.get("response", "")) state["conversation"].append({"role": "user", "content": user_msg}) state["full_log"].append(parsed) enforce = parsed.get("enforce_preferences", False) if isinstance(enforce, str): enforce = enforce.lower() == "true" if enforce: state["enforcement_count"] += 1 # Detect disappointment and satisfaction from user message # Disappointment indicators (not quite right, could be better, etc.) user_msg_lower = user_msg.lower() disappointment = any(phrase in user_msg_lower for phrase in [ "not quite", "not what i", "that's not", "incorrect", "wrong", "mistake", "error", "confused", "doesn't make sense", "try again", "not helpful", "not useful" ]) # Satisfaction indicators (explicit positive feedback) satisfaction = parsed.get("should_terminate", False) or any(phrase in user_msg_lower for phrase in [ "perfect", "exactly", "great", "thanks", "helpful", "that's right", "correct", "good job", "well done", "makes sense", "understand now", "got it" ]) # Store parsed feedback for REINFORCE (applied AFTER prepare_prompt sets pending_rl_update) state["_pending_feedback"] = { "user_msg": user_msg, "enforce": bool(enforce), "disappointment": disappointment and not enforce, # Don't double-count "satisfaction": satisfaction and not enforce, # Don't count if also enforcing "draft_answer": bool(parsed.get("draft_answer")), } if parsed.get("should_terminate", False) or TERMINATION_SIGNAL in user_msg: to_remove.append(pidx) continue # Batch preference extraction for PersonalizedLLM adapters extraction_batch = [] # (pidx, query) remaining_active = [pidx for pidx in active_list if pidx not in to_remove] for pidx in remaining_active: adapter = adapters.get(pidx) if adapter and hasattr(adapter, '_llm') and hasattr(adapter._llm, 'enable_preference_extraction'): if adapter._llm.enable_preference_extraction and adapter._llm._extractor is not None: query = adapter._llm.get_last_user_query(adapter._current_user_id) if hasattr(adapter._llm, 'get_last_user_query') else None if not query: state = all_states[pidx] query = state["conversation"][-1]["content"] if state["conversation"] else "" if query: extraction_batch.append((pidx, query)) if extraction_batch: extractor = extraction_batch[0][1] # just need any adapter to get the extractor adapter0 = adapters[extraction_batch[0][0]] shared_extractor = adapter0._llm._extractor if hasattr(shared_extractor, 'batch_extract_preferences'): queries = [q for _, q in extraction_batch] batch_results = shared_extractor.batch_extract_preferences(queries) for (pidx, _), pref_dict in zip(extraction_batch, batch_results): adapter = adapters[pidx] adapter._llm.apply_extracted_preferences(adapter._current_user_id, pref_dict) else: # Fallback: sequential for pidx, query in extraction_batch: adapter = adapters[pidx] adapter._llm._extractor.extract_turn(adapter._llm._sessions[adapter._current_user_id].session_state.history) # Batch scaffolding for reflection adapters before prepare_prompt scaffolding_batch = [] # (pidx, prompt) remaining_active = [pidx for pidx in active_list if pidx not in to_remove] for pidx in remaining_active: adapter = adapters.get(pidx) if adapter and hasattr(adapter, 'get_scaffolding_prompt'): state = all_states[pidx] # Temporarily add user msg to history for scaffolding agent_notes = adapter._user_notes.get(adapter._current_user_id, "No notes yet about this user.") if adapter.with_scaffolding and agent_notes != "No notes yet about this user.": prompt = adapter.get_scaffolding_prompt( state["conversation"], agent_notes) if prompt is not None: scaffolding_batch.append((pidx, prompt)) if scaffolding_batch: scaff_messages = [[{"role": "user", "content": p}] for _, p in scaffolding_batch] scaff_responses = agent_client.batch_completion(scaff_messages) for (pidx, _), resp in zip(scaffolding_batch, scaff_responses): adapter = adapters[pidx] adapter._scaffolding_result = resp if resp else None # Prepare agent prompts for batching # NOTE: prepare_prompt calls chat_prepare which sets pending_rl_update # from the previous turn's data. REINFORCE feedback must be applied # AFTER this call so that pending_rl_update is available. for pidx in remaining_active: state = all_states[pidx] try: adapter = adapters[pidx] user_msg = state["conversation"][-1]["content"] if hasattr(adapter, 'prepare_prompt'): messages, context = adapter.prepare_prompt(user_msg, state["conversation"][:-1]) agent_prompts_batch.append((pidx, messages, context)) elif hasattr(adapter, 'generate_response'): agent_prompts_batch.append((pidx, None, None)) else: state["conversation"].append({"role": "assistant", "content": "[Error: Adapter not configured]"}) except Exception as e: logger.error(f" Agent prepare error p{pidx} t{turn}: {e}") state["conversation"].append({"role": "assistant", "content": "I apologize, I encountered an error. Could you rephrase?"}) # Apply REINFORCE feedback NOW (after prepare_prompt set pending_rl_update) for pidx in remaining_active: state = all_states[pidx] fb = state.pop("_pending_feedback", None) if fb: adapter = adapters.get(pidx) if adapter and hasattr(adapter, 'process_user_turn'): adapter.process_user_turn( user_response=fb["user_msg"], enforce_preferences=fb["enforce"], express_disappointment=fb.get("disappointment", False), express_satisfaction=fb["satisfaction"], draft_answer_updated=fb["draft_answer"], ) # Also apply feedback for terminated sessions (they skipped prepare_prompt # but still need the reward signal from their last turn) for pidx in to_remove: state = all_states.get(pidx) if not state: continue fb = state.pop("_pending_feedback", None) if fb: adapter = adapters.get(pidx) if adapter and hasattr(adapter, 'process_user_turn'): # For terminated sessions, we can't call prepare_prompt # (no next turn), but we still want the reward applied. # Call chat_prepare with a dummy to set pending_rl_update, # then apply feedback. try: if hasattr(adapter, '_llm') and hasattr(adapter._llm, 'chat_prepare'): adapter._llm.chat_prepare( adapter._current_user_id, fb["user_msg"], skip_extraction=True, skip_auto_reward=True, ) adapter.process_user_turn( user_response=fb["user_msg"], enforce_preferences=fb["enforce"], express_disappointment=fb.get("disappointment", False), express_satisfaction=fb["satisfaction"], draft_answer_updated=fb["draft_answer"], ) except Exception: pass # Best effort for terminated sessions # Batch vLLM call for all agent prompts if agent_prompts_batch: # Separate prompts that can be batched from fallback batchable = [(pidx, msgs, ctx) for pidx, msgs, ctx in agent_prompts_batch if msgs is not None] fallback = [(pidx, msgs, ctx) for pidx, msgs, ctx in agent_prompts_batch if msgs is None] # Batch call for batchable prompts if batchable: batch_messages = [msgs for _, msgs, _ in batchable] batch_responses = agent_client.batch_completion(batch_messages) # Process batched responses for (pidx, _, context), response in zip(batchable, batch_responses): try: adapter = adapters[pidx] state = all_states[pidx] if response is not None: result = adapter.process_response(response, context) agent_content = result.get("response", str(result)) if isinstance(result, dict) else str(result) else: agent_content = "I apologize, I encountered an error. Could you rephrase?" state["conversation"].append({"role": "assistant", "content": agent_content}) except Exception as e: logger.error(f" Agent process error p{pidx} t{turn}: {e}") all_states[pidx]["conversation"].append({"role": "assistant", "content": "I apologize, I encountered an error. Could you rephrase?"}) # Handle fallback (adapters without prepare_prompt - sequential calls) for pidx, _, _ in fallback: try: adapter = adapters[pidx] state = all_states[pidx] user_msg = state["conversation"][-1]["content"] resp = adapter.generate_response(user_msg, state["conversation"][:-1]) agent_content = resp.get("response", str(resp)) if isinstance(resp, dict) else str(resp) state["conversation"].append({"role": "assistant", "content": agent_content}) except Exception as e: logger.error(f" Agent fallback error p{pidx} t{turn}: {e}") all_states[pidx]["conversation"].append({"role": "assistant", "content": "I apologize, I encountered an error. Could you rephrase?"}) active_set -= set(to_remove) # Batch note-update for reflection adapters before end_session note_update_batch = [] # (profile_idx, messages) for profile_idx in profiles_to_run: if profile_idx not in all_states: continue adapter = adapters.get(profile_idx) if adapter and hasattr(adapter, 'get_note_update_prompt'): prompt_msgs = adapter.get_note_update_prompt() if prompt_msgs is not None: note_update_batch.append((profile_idx, prompt_msgs)) if note_update_batch: note_messages = [msgs for _, msgs in note_update_batch] note_responses = agent_client.batch_completion(note_messages) for (profile_idx, _), resp in zip(note_update_batch, note_responses): if resp: adapter = adapters[profile_idx] adapter.apply_note_update_response(resp) # Save results for this session round for profile_idx in profiles_to_run: if profile_idx not in all_states: continue state = all_states[profile_idx] problem_dict = state["problem_dict"] conversation = state["conversation"] full_log = state["full_log"] user_tokens = sum(len(m["content"].split()) for m in conversation if m["role"] == "user") agent_tokens = sum(len(m["content"].split()) for m in conversation if m["role"] == "assistant") enforcement_count = state["enforcement_count"] task_success = 0 for entry in full_log: if entry.get("should_terminate", False): draft = str(entry.get("draft_answer", "")) if draft and "don't know" not in draft.lower(): task_success = 1 # End session on adapter (applies task completion reward for REINFORCE) adapter = adapters.get(profile_idx) if adapter and hasattr(adapter, 'end_session'): # Skip note update if batch already handled it skip_notes = hasattr(adapter, 'get_note_update_prompt') try: adapter.end_session(task_success=bool(task_success), skip_note_update=skip_notes) except TypeError: adapter.end_session(task_success=bool(task_success)) results.append({ "method": method, "profile_id": self.profiles[profile_idx].get("user_id", f"user_{profile_idx}"), "problem_id": f"session_{session_idx}", "problem": problem_dict.get("problem", ""), "ground_truth_solution": problem_dict.get("solution", ""), "is_conflict_test": state["is_conflict"], "conversation": {"turns": conversation}, "full_user_log": full_log, "metrics": { "task_success": bool(task_success), "total_turns": len(conversation), "user_token_count": user_tokens, "agent_token_count": agent_tokens, "total_token_count": user_tokens + agent_tokens, "enforcement_count": enforcement_count, "disappointment_count": 0, "preference_compliance_scores": [], "conflict_resolution_accuracy": 0, "over_personalization_rate": 0, }, "adapter_metrics": {}, }) # Collect adapter metrics (e.g. user_vector_norm for rag_vector) adapter = adapters.get(profile_idx) if adapter and hasattr(adapter, 'get_user_vector'): user_id = self.profiles[profile_idx].get("user_id", f"user_{profile_idx}") vec = adapter.get_user_vector(user_id) if vec is not None: results[-1]["adapter_metrics"] = { "user_vector_norm": float(np.linalg.norm(vec)), } # Save user vector snapshots every 10 sessions if (session_idx + 1) % 10 == 0: vectors_dir = checkpoint_file.parent / "vectors" vectors_dir.mkdir(parents=True, exist_ok=True) user_vectors = {} for profile_idx in profiles_to_run: adapter = adapters.get(profile_idx) if adapter and hasattr(adapter, 'get_user_vector'): user_id = self.profiles[profile_idx].get("user_id", f"user_{profile_idx}") vec = adapter.get_user_vector(user_id) if vec is not None: user_vectors[user_id] = vec if user_vectors: snapshot_path = vectors_dir / f"vectors_session_{session_idx+1}.npy" np.save(snapshot_path, user_vectors) logger.info(f" Saved {len(user_vectors)} user vectors to {snapshot_path}") # Checkpoint after each session round with session-level tracking # Only increment for profiles that actually ran in this round (those in all_states) for profile_idx in all_states.keys(): sessions_per_profile[str(profile_idx)] = sessions_per_profile.get(str(profile_idx), 0) + 1 if sessions_per_profile[str(profile_idx)] >= self.config.n_sessions_per_profile: completed_profiles.add(profile_idx) with open(checkpoint_file, "w") as f: json.dump({ "completed_profiles": sorted(list(completed_profiles)), "sessions_per_profile": sessions_per_profile }, f) with open(results_file, "w") as f: json.dump(results, f, indent=2) elapsed = time.time() - start_time sessions_done = len(results) rate = sessions_done / elapsed * 3600 if elapsed > 0 else 0 logger.info(f" Session round {session_idx+1}/{n_sessions}: {sessions_done} total, {rate:.0f} sessions/hr") # Export user vectors before cleanup (for RAG methods with user vectors) self._export_user_vectors(method, adapters) # Explicitly free adapter models to prevent GPU OOM across methods for pidx, adapter in adapters.items(): if hasattr(adapter, 'cleanup'): adapter.cleanup() del adapters return results def run_all(self) -> Dict[str, Any]: """Run all methods and generate comparative analysis.""" all_results = {} for method in self.config.methods: if method not in AVAILABLE_METHODS: logger.warning(f"Unknown method: {method}, skipping") continue results = self.run_method(method) all_results[method] = results # Free GPU memory between methods to prevent OOM on later adapters try: from personalization.serving.personalized_llm import clear_shared_models clear_shared_models() except ImportError: pass try: import gc import torch gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() logger.info(f" GPU memory freed after {method}: {torch.cuda.memory_allocated()/1e9:.1f}GB allocated") except ImportError: pass # Comparative analysis analysis = self._analyze_results(all_results) # Save analysis with open(self.output_dir / "analysis.json", "w") as f: json.dump(analysis, f, indent=2) # Generate report self._generate_report(analysis) return analysis def _analyze_results(self, all_results: Dict[str, List[Dict]]) -> Dict: """Analyze results across all methods.""" analysis = { "per_method": {}, "comparison": {}, } for method, results in all_results.items(): n = len(results) if n == 0: continue # Aggregate metrics task_success = sum(r["metrics"]["task_success"] for r in results) / n avg_user_tokens = sum(r["metrics"]["user_token_count"] for r in results) / n avg_total_tokens = sum(r["metrics"]["total_token_count"] for r in results) / n avg_enforcement = sum(r["metrics"]["enforcement_count"] for r in results) / n avg_turns = sum(r["metrics"]["total_turns"] for r in results) / n # Compliance and conflict metrics compliance_scores = [ sum(r["metrics"]["preference_compliance_scores"]) / len(r["metrics"]["preference_compliance_scores"]) if r["metrics"]["preference_compliance_scores"] else 0.5 for r in results ] avg_compliance = sum(compliance_scores) / len(compliance_scores) conflict_results = [r for r in results if r["is_conflict_test"]] conflict_accuracy = sum( r["metrics"]["conflict_resolution_accuracy"] for r in conflict_results ) / len(conflict_results) if conflict_results else 0 over_personalization = sum( r["metrics"]["over_personalization_rate"] for r in results ) / n analysis["per_method"][method] = { "n_sessions": n, "task_success_rate": task_success, "avg_user_tokens": avg_user_tokens, "avg_total_tokens": avg_total_tokens, "avg_enforcement_count": avg_enforcement, "avg_turns": avg_turns, "avg_preference_compliance": avg_compliance, "conflict_resolution_accuracy": conflict_accuracy, "over_personalization_rate": over_personalization, } # Comparison metrics_to_compare = [ ("task_success_rate", True), # higher is better ("avg_user_tokens", False), # lower is better ("avg_total_tokens", False), # lower is better ("avg_enforcement_count", False), # lower is better ("avg_preference_compliance", True), # higher is better ("conflict_resolution_accuracy", True), # higher is better ("over_personalization_rate", False), # lower is better ] for metric, higher_better in metrics_to_compare: values = {m: analysis["per_method"][m][metric] for m in analysis["per_method"]} if not values: logger.warning(f"No values for metric {metric}, skipping comparison") continue if higher_better: best = max(values, key=values.get) else: best = min(values, key=values.get) analysis["comparison"][metric] = { "values": values, "best_method": best, "best_value": values[best], } return analysis def _generate_report(self, analysis: Dict) -> None: """Generate a human-readable report.""" report_lines = [ "# Personalization Experiment Report", f"\nGenerated: {datetime.now().isoformat()}", f"\nConfig: {self.config.n_profiles} profiles, {self.config.n_sessions_per_profile} sessions each", "\n## Method Comparison\n", ] # Create comparison table metrics_display = [ ("Task Success", "task_success_rate", "{:.1%}"), ("User Effort (tokens)", "avg_user_tokens", "{:.0f}"), ("Total Tokens", "avg_total_tokens", "{:.0f}"), ("Enforcement Count", "avg_enforcement_count", "{:.2f}"), ("Preference Compliance", "avg_preference_compliance", "{:.1%}"), ("Conflict Resolution", "conflict_resolution_accuracy", "{:.1%}"), ("Over-personalization", "over_personalization_rate", "{:.1%}"), ] methods = list(analysis["per_method"].keys()) # Header header = "| Metric |" + "|".join(f" {m} " for m in methods) + "| Best |" separator = "|" + "|".join(["-" * (len(m) + 2) for m in ["Metric"] + methods + ["Best"]]) + "|" report_lines.extend([header, separator]) for display_name, metric_key, fmt in metrics_display: row = f"| {display_name} |" for m in methods: val = analysis["per_method"].get(m, {}).get(metric_key, 0) row += f" {fmt.format(val)} |" if metric_key in analysis.get("comparison", {}): best = analysis["comparison"][metric_key]["best_method"] else: best = "N/A" row += f" {best} |" report_lines.append(row) # Key findings report_lines.extend([ "\n## Key Findings\n", ]) # Find advantages of proposed methods rag_vector = analysis["per_method"].get("rag_vector", {}) rag = analysis["per_method"].get("rag", {}) contextual = analysis["per_method"].get("contextual", {}) all_memory = analysis["per_method"].get("all_memory", {}) if rag_vector and contextual: token_reduction = (contextual.get("avg_total_tokens", 0) - rag_vector.get("avg_total_tokens", 0)) / contextual.get("avg_total_tokens", 1) * 100 report_lines.append(f"- **Token Efficiency**: RAG+Vector uses {token_reduction:.1f}% fewer tokens than contextual memory") if rag_vector and all_memory: conflict_improvement = rag_vector.get("conflict_resolution_accuracy", 0) - all_memory.get("conflict_resolution_accuracy", 0) report_lines.append(f"- **Conflict Resolution**: RAG+Vector improves by {conflict_improvement:.1%} over all-memory baseline") if rag_vector: report_lines.append(f"- **Over-personalization**: RAG+Vector rate: {rag_vector.get('over_personalization_rate', 0):.1%}") # Save report report_path = self.output_dir / "report.md" with open(report_path, "w") as f: f.write("\n".join(report_lines)) logger.info(f"Report saved to {report_path}") def main(): parser = argparse.ArgumentParser(description="Run personalization experiments") parser.add_argument("--config", type=str, help="Path to config YAML file") parser.add_argument("--methods", type=str, default="vanilla,contextual,rag,rag_vector", help="Comma-separated list of methods to compare") parser.add_argument("--datasets", type=str, default="math-hard,math-500,bigcodebench", help="Comma-separated list of datasets") parser.add_argument("--n-profiles", type=int, default=200, help="Number of user profiles") parser.add_argument("--n-sessions", type=int, default=30, help="Sessions per profile") parser.add_argument("--max-turns", type=int, default=15, help="Max turns per session") parser.add_argument("--output-dir", type=str, default="results", help="Output directory") parser.add_argument("--profile-path", type=str, help="Path to pre-generated profiles") parser.add_argument("--start-profile", type=int, default=0, help="Start profile index (inclusive, 0-indexed)") parser.add_argument("--end-profile", type=int, default=None, help="End profile index (exclusive). If not set, runs all profiles from start") # vLLM and parallel processing options parser.add_argument("--use-vllm", action="store_true", help="Use vLLM servers for inference (much faster)") parser.add_argument("--vllm-user-url", type=str, default="http://localhost:8004/v1", help="vLLM server URL for user simulator (70B)") parser.add_argument("--vllm-agent-url", type=str, default="http://localhost:8003/v1", help="vLLM server URL for agent (8B)") # OpenAI user agent options parser.add_argument("--use-openai-user", action="store_true", help="Use OpenAI API (GPT-5) for user simulation instead of vLLM") parser.add_argument("--openai-user-model", type=str, default="gpt-5", help="OpenAI model name for user simulator (default: gpt-5)") parser.add_argument("--reward-mode", type=str, default="keyword", choices=["keyword", "llm", "llm_local"], help="Reward mode: 'keyword' (user signals), 'llm' (GPT-4o-mini), or 'llm_local' (local vLLM)") parser.add_argument("--reward-vllm-url", type=str, default="http://localhost:8005/v1", help="vLLM server URL for local reward model (when --reward-mode=llm_local)") parser.add_argument("--parallel-profiles", type=int, default=50, help="Number of profiles to process in parallel (requires --use-vllm)") parser.add_argument("--use-batch-processing", action="store_true", default=True, help="Use turn-synchronous batch processing for vanilla/all_memory") parser.add_argument("--no-batch-processing", action="store_false", dest="use_batch_processing", help="Disable batch processing") parser.add_argument("--batch-size", type=int, default=50, help="Number of conversations to batch together") parser.add_argument("--continue-from", type=str, default=None, help="Path to existing output directory to continue from (for extending sessions)") args = parser.parse_args() # Load or create config if args.config and Path(args.config).exists(): with open(args.config) as f: config_dict = yaml.safe_load(f) config = ExperimentConfig(**config_dict) else: config = ExperimentConfig( methods=args.methods.split(","), datasets=args.datasets.split(","), n_profiles=args.n_profiles, n_sessions_per_profile=args.n_sessions, max_turns_per_session=args.max_turns, output_dir=args.output_dir, profile_path=args.profile_path, start_profile=args.start_profile, end_profile=args.end_profile, use_vllm=args.use_vllm, vllm_user_url=args.vllm_user_url, vllm_agent_url=args.vllm_agent_url, use_openai_user=args.use_openai_user, openai_user_model=args.openai_user_model, reward_mode=args.reward_mode, reward_vllm_url=args.reward_vllm_url, parallel_profiles=args.parallel_profiles, use_batch_processing=args.use_batch_processing, batch_size_conversations=args.batch_size, continue_from=args.continue_from, ) # Run experiments runner = ExperimentRunner(config) analysis = runner.run_all() print("\n" + "=" * 60) print("EXPERIMENT COMPLETE") print("=" * 60) print(f"\nResults saved to: {runner.output_dir}") if analysis.get("comparison"): print("\nBest methods per metric:") for metric, data in analysis["comparison"].items(): print(f" {metric}: {data['best_method']} ({data['best_value']:.3f})") else: print("\nNo comparison data available (sessions may have failed)") if __name__ == "__main__": main()