path: root/collaborativeagents/agents/batch_vllm_agent.py

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