Add collaborativeagents module and update gitignore

- Add collaborativeagents subproject with adapters, agents, and evaluation modules
- Update .gitignore to exclude large binary files (.whl, .tar), wandb logs, and results

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

1 files changed, 643 insertions, 0 deletions

diff --git a/collaborativeagents/agents/batch_vllm_agent.py b/collaborativeagents/agents/batch_vllm_agent.py
new file mode 100644
index 0000000..f57bd78
--- /dev/null
+++ b/collaborativeagents/agents/batch_vllm_agent.py
@@ -0,0 +1,643 @@
+"""
+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)