""" Batch processing for high-throughput conversation generation. This implements turn-synchronous batch processing: - vLLM servers for local models (agent) - OpenAI async SDK for API-based models (user simulator) Key insight: Process ALL conversations at the same turn level together, maximizing throughput via concurrent async requests. """ import asyncio import aiohttp import os from typing import List, Dict, Any, Optional from copy import deepcopy from json_repair import repair_json import time TERMINATION_SIGNAL = "TERMINATE" class BatchVLLMClient: """ Async batch client for vLLM server. Sends multiple requests concurrently and gathers results, allowing vLLM's continuous batching to process them together. """ def __init__( self, vllm_url: str, model_name: str = None, max_tokens: int = 512, temperature: float = 0.7, timeout: int = None, # None = infinite timeout max_concurrent: int = 100, api_key: str = None, is_reasoning_model: bool = False, json_mode: bool = False, # Enable JSON output mode for vLLM ): self.vllm_url = vllm_url.rstrip('/') self.model_name = model_name self.max_tokens = max_tokens self.temperature = temperature self.timeout = timeout # None for infinite self.max_concurrent = max_concurrent self.api_key = api_key self.is_reasoning_model = is_reasoning_model self.json_mode = json_mode if self.model_name is None and not self.api_key: self._discover_model_sync() def _discover_model_sync(self): """Synchronously discover model name.""" import requests try: response = requests.get(f"{self.vllm_url}/models", timeout=10) response.raise_for_status() models = response.json() if models.get("data") and len(models["data"]) > 0: self.model_name = models["data"][0]["id"] else: self.model_name = "default" except Exception as e: print(f"[BatchVLLMClient] Warning: Could not discover model ({e})") self.model_name = "default" async def _single_completion( self, session: aiohttp.ClientSession, messages: List[Dict[str, str]], idx: int, retry_with_lower_tokens: bool = True ) -> tuple: """Make a single async completion request with retry logic.""" max_tokens = self.max_tokens for attempt in range(3): # Up to 3 attempts if self.is_reasoning_model: payload = { "model": self.model_name, "messages": messages, "max_completion_tokens": max_tokens, "response_format": {"type": "json_object"}, } elif self.api_key: payload = { "model": self.model_name, "messages": messages, "max_tokens": max_tokens, "temperature": self.temperature, "response_format": {"type": "json_object"}, } else: payload = { "model": self.model_name, "messages": messages, "max_tokens": max_tokens, "temperature": self.temperature, "top_p": 0.9, } # Add JSON mode for local vLLM if requested if self.json_mode: payload["response_format"] = {"type": "json_object"} try: # Use None for infinite timeout timeout_config = aiohttp.ClientTimeout(total=self.timeout) if self.timeout else None async with session.post( f"{self.vllm_url}/chat/completions", json=payload, timeout=timeout_config ) as response: if response.status == 200: result = await response.json() content = result["choices"][0]["message"]["content"] # Reasoning models may exhaust tokens on internal reasoning if not content and self.is_reasoning_model: finish = result["choices"][0].get("finish_reason", "") if finish == "length": max_tokens = min(max_tokens * 2, 8192) continue return (idx, content, None) elif response.status == 400: error_text = await response.text() if "max_tokens" in error_text and retry_with_lower_tokens: max_tokens = max(64, max_tokens // 2) continue return (idx, None, f"HTTP 400: {error_text[:200]}") elif response.status == 429: # Rate limit — wait and retry await asyncio.sleep(2 ** attempt) continue else: error_text = await response.text() return (idx, None, f"HTTP {response.status}: {error_text[:200]}") except asyncio.TimeoutError: if attempt < 2: continue return (idx, None, "Timeout") except Exception as e: return (idx, None, str(e)) return (idx, None, "Max retries exceeded") async def batch_completion_async( self, messages_list: List[List[Dict[str, str]]], show_progress: bool = False ) -> List[Optional[str]]: """ Send multiple completion requests concurrently. vLLM's continuous batching will automatically batch these together. Uses semaphore to limit concurrent requests. """ results = [None] * len(messages_list) errors = [None] * len(messages_list) completed = 0 # Use semaphore to limit concurrent requests semaphore = asyncio.Semaphore(self.max_concurrent) async def limited_completion(session, messages, idx): async with semaphore: return await self._single_completion(session, messages, idx) connector = aiohttp.TCPConnector(limit=self.max_concurrent) headers = {"Content-Type": "application/json"} if self.api_key: headers["Authorization"] = f"Bearer {self.api_key}" async with aiohttp.ClientSession(connector=connector, headers=headers) as session: tasks = [ limited_completion(session, messages, idx) for idx, messages in enumerate(messages_list) ] for coro in asyncio.as_completed(tasks): idx, content, error = await coro completed += 1 if error: errors[idx] = error results[idx] = content if show_progress and completed % 10 == 0: print(f" [{completed}/{len(messages_list)}] completed") # Summary of errors error_count = sum(1 for e in errors if e is not None) if error_count > 0: print(f"[BatchVLLMClient] {error_count}/{len(messages_list)} requests failed") return results def batch_completion( self, messages_list: List[List[Dict[str, str]]] ) -> List[Optional[str]]: """Synchronous wrapper for batch completion.""" return asyncio.run(self.batch_completion_async(messages_list)) class BatchOpenAIClient: """ Async batch client using the OpenAI Python SDK. Drop-in replacement for BatchVLLMClient when targeting OpenAI API. Uses AsyncOpenAI for proper SSL, auth, and reasoning model support. """ def __init__( self, model: str = "gpt-5", max_tokens: int = 4096, max_concurrent: int = 32, timeout: float = 120.0, api_key: str = None, ): from openai import AsyncOpenAI self.model = model self.max_tokens = max_tokens self.max_concurrent = max_concurrent self.timeout = timeout self._is_reasoning = any(model.startswith(p) for p in ("o1", "o3", "gpt-5")) self._client = AsyncOpenAI( api_key=api_key or os.environ.get("OPENAI_API_KEY"), timeout=timeout, max_retries=2, ) async def _single_completion( self, messages: List[Dict[str, str]], idx: int, ) -> tuple: """Single async completion via OpenAI SDK.""" max_tokens = self.max_tokens for attempt in range(3): try: kwargs = { "model": self.model, "messages": messages, "response_format": {"type": "json_object"}, } if self._is_reasoning: kwargs["max_completion_tokens"] = max_tokens else: kwargs["max_tokens"] = max_tokens kwargs["temperature"] = 0.7 response = await self._client.chat.completions.create(**kwargs) content = response.choices[0].message.content # Reasoning models may exhaust tokens on internal reasoning if not content and self._is_reasoning: if response.choices[0].finish_reason == "length": max_tokens = min(max_tokens * 2, 16384) continue return (idx, content, None) except Exception as e: if attempt < 2: await asyncio.sleep(1 * (attempt + 1)) continue return (idx, None, f"{type(e).__name__}: {e}") return (idx, None, "Max retries exceeded") async def batch_completion_async( self, messages_list: List[List[Dict[str, str]]], show_progress: bool = False ) -> List[Optional[str]]: """Send multiple completion requests concurrently via AsyncOpenAI.""" results = [None] * len(messages_list) errors = [None] * len(messages_list) completed = 0 semaphore = asyncio.Semaphore(self.max_concurrent) async def limited_completion(messages, idx): async with semaphore: return await self._single_completion(messages, idx) tasks = [ limited_completion(messages, idx) for idx, messages in enumerate(messages_list) ] for coro in asyncio.as_completed(tasks): idx, content, error = await coro completed += 1 if error: errors[idx] = error results[idx] = content if show_progress and completed % 10 == 0: print(f" [BatchOpenAI {completed}/{len(messages_list)}] completed") error_count = sum(1 for e in errors if e is not None) if error_count > 0: error_samples = [e for e in errors if e is not None][:3] print(f"[BatchOpenAIClient] {error_count}/{len(messages_list)} requests failed. Examples: {error_samples}") return results def batch_completion( self, messages_list: List[List[Dict[str, str]]] ) -> List[Optional[str]]: """Synchronous wrapper for batch completion.""" return asyncio.run(self.batch_completion_async(messages_list)) class BatchConversationGenerator: """ Generate conversations using turn-synchronous batch processing. This processes ALL samples at the same turn together, maximizing vLLM's continuous batching efficiency. """ USER_SYSTEM_PROMPT = """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. # Problem Description {problem} Note: the agent cannot see this problem description. # User Persona {user_persona} # User Preferences {user_preferences} These preferences are NON-NEGOTIABLE that define how you prefer the agent to behave. They must be strictly enforced: - **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. - **Enforce immediately**: Every agent response must satisfy your preferences before you can proceed. Explicitly ask the agent to adjust their response until it complies. - **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. # Draft Answer Management - **Maintain a working draft**: Start with "I don't know". Update your draft answer based on what you learn from agent responses. - **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. # Conversation Termination Before generating your response, determine if you should terminate: - Do you feel like your draft answer is a good answer to the problem? - Do you feel like the agent cannot help further? If the agent response does not follow your preferences, you must NOT terminate - instead, enforce the preferences. When ready to terminate, respond with "TERMINATE". # Output Format: {{ "preferences_check": "For EACH relevant preference, evaluate: is it satisfied?", "enforce_preferences": true/false, "reasoning": "Brief reasoning (2-3 sentences). Does agent follow preferences? If no, enforce. If yes, update draft.", "draft_answer": "Your current working draft answer", "should_terminate": true/false, "response": "Your response to the agent" }} """ USER_SYSTEM_PROMPT_NO_PREF = """You are a user simulator collaborating with an agent to solve a problem. # Problem Description {problem} # User Persona {user_persona} # Output Format: {{ "reasoning": "Brief reasoning", "draft_answer": "Your current draft answer", "should_terminate": true/false, "response": "Your response" }} """ def __init__( self, user_vllm_url: str, agent_vllm_url: str, max_turns: int = 10, user_max_tokens: int = 512, agent_max_tokens: int = 1024, temperature: float = 0.7, user_api_key: str = None, user_model_name: str = None, user_is_reasoning: bool = False, ): if user_api_key: # Use OpenAI SDK client for API-based user models self.user_client = BatchOpenAIClient( model=user_model_name or "gpt-5", max_tokens=user_max_tokens, max_concurrent=32, timeout=120.0, api_key=user_api_key, ) else: self.user_client = BatchVLLMClient( vllm_url=user_vllm_url, model_name=user_model_name, max_tokens=user_max_tokens, temperature=temperature, ) self.agent_client = BatchVLLMClient( vllm_url=agent_vllm_url, max_tokens=agent_max_tokens, temperature=temperature, ) self.max_turns = max_turns def _build_user_system_prompt( self, problem: str, user_persona: str, user_preferences: str = None ) -> str: if user_preferences: return self.USER_SYSTEM_PROMPT.format( problem=problem, user_persona=user_persona, user_preferences=user_preferences ) else: return self.USER_SYSTEM_PROMPT_NO_PREF.format( problem=problem, user_persona=user_persona ) def _reverse_roles(self, conversation: List[Dict]) -> List[Dict]: """Reverse roles for user perspective.""" return [ {"role": "user" if msg["role"] == "assistant" else "assistant", "content": msg["content"]} for msg in conversation ] def _parse_user_response(self, content: str) -> Optional[Dict]: """Parse user response JSON.""" if content is None: return None try: parsed = repair_json(content, return_objects=True) required_keys = ["reasoning", "draft_answer", "should_terminate", "response"] if all(k in parsed for k in required_keys): return parsed # Fallback: treat as raw response if TERMINATION_SIGNAL in content: return { "reasoning": "", "draft_answer": "", "should_terminate": True, "response": TERMINATION_SIGNAL } return { "reasoning": "", "draft_answer": "", "should_terminate": False, "response": content } except: if TERMINATION_SIGNAL in content: return { "reasoning": "", "draft_answer": "", "should_terminate": True, "response": TERMINATION_SIGNAL } return { "reasoning": "", "draft_answer": "", "should_terminate": False, "response": content } def generate_batch( self, samples: List[Dict], user_persona: str, user_preferences: str = None, agent_system_prompt: str = "You are a helpful AI assistant.", ) -> List[Dict]: """ Generate conversations for a batch of samples using turn-synchronous processing. Args: samples: List of dicts with 'problem' key user_persona: User persona description user_preferences: User preferences string (optional) agent_system_prompt: System prompt for the agent Returns: List of conversation results """ n_samples = len(samples) # Initialize state for all conversations conversations = [[{"role": "assistant", "content": "How can I help you?"}] for _ in range(n_samples)] full_logs = [[] for _ in range(n_samples)] # Build user system prompts for each sample user_system_prompts = [ self._build_user_system_prompt( problem=sample['problem'], user_persona=user_persona, user_preferences=user_preferences ) for sample in samples ] # Track active conversations (not terminated or failed) active_indices = set(range(n_samples)) failed_indices = set() for turn in range(self.max_turns): if not active_indices: break # ========== USER TURN (BATCHED) ========== user_indices = sorted(active_indices) # Build batch of user messages user_messages_batch = [] for idx in user_indices: messages = [{"role": "system", "content": user_system_prompts[idx]}] messages.extend(self._reverse_roles(conversations[idx])) user_messages_batch.append(messages) # Batch call to user model user_responses_raw = self.user_client.batch_completion(user_messages_batch) # Process user responses for i, idx in enumerate(user_indices): raw = user_responses_raw[i] parsed = self._parse_user_response(raw) if parsed is None: active_indices.discard(idx) failed_indices.add(idx) continue conversations[idx].append({ "role": "user", "content": str(parsed["response"]) }) full_logs[idx].append(parsed) # Check for termination if parsed.get("should_terminate") or TERMINATION_SIGNAL in parsed["response"]: active_indices.discard(idx) if not active_indices: break # ========== AGENT TURN (BATCHED) ========== agent_indices = sorted(active_indices) # Build batch of agent messages agent_messages_batch = [] for idx in agent_indices: messages = [{"role": "system", "content": agent_system_prompt}] messages.extend(conversations[idx]) agent_messages_batch.append(messages) # Batch call to agent model agent_responses_raw = self.agent_client.batch_completion(agent_messages_batch) # Process agent responses for i, idx in enumerate(agent_indices): raw = agent_responses_raw[i] if raw is None: active_indices.discard(idx) failed_indices.add(idx) continue conversations[idx].append({ "role": "assistant", "content": raw }) full_logs[idx].append({"response": raw, "reasoning": ""}) # Build results results = [] for i, sample in enumerate(samples): if i in failed_indices: results.append(None) else: results.append({ "sample": sample, "conversation": conversations[i], "full_conversation_log": full_logs[i] }) return results def benchmark_batch_generation( user_url: str, agent_url: str, n_samples: int = 20, max_turns: int = 5, ): """Quick benchmark of batch generation.""" # Create dummy samples samples = [ {"problem": f"What is {i+1} + {i+2}? Show your work.", "solution": str(2*i+3)} for i in range(n_samples) ] generator = BatchConversationGenerator( user_vllm_url=user_url, agent_vllm_url=agent_url, max_turns=max_turns, ) start = time.time() results = generator.generate_batch( samples=samples, user_persona="A curious student learning math.", user_preferences="1. Show step-by-step working\n2. Explain clearly", ) elapsed = time.time() - start successes = sum(1 for r in results if r is not None) print(f"\n=== Batch Generation Benchmark ===") print(f"Samples: {n_samples}, Max turns: {max_turns}") print(f"Successes: {successes}/{n_samples}") print(f"Time: {elapsed:.1f}s") print(f"Throughput: {successes * 3600 / elapsed:.0f} conversations/hr") return results if __name__ == "__main__": import sys if len(sys.argv) >= 3: user_url = sys.argv[1] agent_url = sys.argv[2] n_samples = int(sys.argv[3]) if len(sys.argv) > 3 else 20 else: user_url = "http://localhost:8004/v1" agent_url = "http://localhost:8003/v1" n_samples = 20 benchmark_batch_generation(user_url, agent_url, n_samples=n_samples)