# Create project directory
mkdir multi-agent-system && cd multi-agent-system

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install httpx asyncio pydantic

# .env
QUBAX_API_KEY=your_api_key_here
QUBAX_BASE_URL=https://api.qubax.ai/v1

import os
import json
import asyncio
from typing import AsyncGenerator
from dataclasses import dataclass
from dotenv import load_dotenv

import httpx

load_dotenv()

@dataclass
class StreamChunk:
    content: str
    role: str = "assistant"

class LLMClient:
    """Streaming LLM client compatible with OpenAI API format."""

    def __init__(self):
        self.api_key = os.getenv("QUBAX_API_KEY")
        self.base_url = os.getenv("QUBAX_BASE_URL", "https://api.qubax.ai/v1")
        self.client = httpx.AsyncClient(
            timeout=httpx.Timeout(300.0, connect=10.0),
            headers={"Authorization": f"Bearer {self.api_key}"}
        )

    async def stream_completion(
        self,
        messages: list[dict],
        model: str = "gpt-5.6-sol",
        temperature: float = 0.7,
        max_tokens: int = 4096
    ) -> AsyncGenerator[StreamChunk, None]:
        """Stream a completion from the LLM, yielding chunks as they arrive."""

        payload = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": True
        }

        async with self.client.stream(
            "POST",
            f"{self.base_url}/chat/completions",
            json=payload,
            headers={"Content-Type": "application/json"}
        ) as response:
            response.raise_for_status()
            async for line in response.aiter_lines():
                if not line or not line.startswith("data: "):
                    continue
                if line.strip() == "data: [DONE]":
                    break
                try:
                    data = json.loads(line[6:])
                    delta = data["choices"][0]["delta"]
                    if "content" in delta and delta["content"]:
                        yield StreamChunk(content=delta["content"])
                except json.JSONDecodeError:
                    continue

    async def complete(
        self,
        messages: list[dict],
        model: str = "gpt-5.6-sol",
        temperature: float = 0.7,
    ) -> str:
        """Non-streaming completion — collects all chunks."""
        full_response = []
        async for chunk in self.stream_completion(messages, model, temperature):
            full_response.append(chunk.content)
        return "".join(full_response)

from enum import Enum
from typing import Optional

class AgentRole(Enum):
    ORCHESTRATOR = "orchestrator"
    RESEARCHER = "researcher"
    CODER = "coder"
    REVIEWER = "reviewer"
    WRITER = "writer"

class Agent:
    """A single AI agent with a specific role and capabilities."""

    def __init__(
        self,
        name: str,
        role: AgentRole,
        model: str = "gpt-5.6-sol",
        system_prompt: Optional[str] = None
    ):
        self.name = name
        self.role = role
        self.model = model
        self.client = LLMClient()
        self.conversation: list[dict] = []

        if system_prompt:
            self.conversation.append({
                "role": "system",
                "content": system_prompt
            })
        else:
            self.conversation.append({
                "role": "system",
                "content": self._default_system_prompt()
            })

    def _default_system_prompt(self) -> str:
        role_descriptions = {
            AgentRole.ORCHESTRATOR: "You are the orchestrator. Break down complex tasks into sub-tasks, delegate to specialist agents, and synthesize their results into a coherent final answer.",
            AgentRole.RESEARCHER: "You are a research agent. Gather information, analyze data, and provide structured findings. Cite sources when possible.",
            AgentRole.CODER: "You are a coding agent. Write clean, well-documented code. Include error handling and follow best practices.",
            AgentRole.REVIEWER: "You are a code reviewer. Check for bugs, security issues, and improvements. Provide specific, actionable feedback.",
            AgentRole.WRITER: "You are a technical writer. Transform technical outputs into clear, readable documentation.",
        }
        return role_descriptions.get(self.role, "You are a helpful AI agent.")

    async def execute(self, task: str, stream: bool = True) -> str:
        """Execute a task and optionally stream the response."""
        self.conversation.append({"role": "user", "content": task})

        if stream:
            print(f"\n[{self.name}] Thinking...", flush=True)
            full_response = []
            async for chunk in self.client.stream_completion(
                self.conversation, model=self.model
            ):
                print(chunk.content, end="", flush=True)
                full_response.append(chunk.content)
            response = "".join(full_response)
        else:
            response = await self.client.complete(
                self.conversation, model=self.model
            )
            print(f"\n[{self.name}] Done.")

        self.conversation.append({"role": "assistant", "content": response})
        return response

import asyncio
from dataclasses import dataclass, field
from typing import Any

@dataclass
class Task:
    id: str
    description: str
    agent_role: AgentRole
    dependencies: list[str] = field(default_factory=list)
    result: str = ""

@dataclass
class TaskPlan:
    tasks: list[Task]

    def get_task(self, task_id: str) -> Optional[Task]:
        for task in self.tasks:
            if task.id == task_id:
                return task
        return None

    def get_ready_tasks(self) -> list[Task]:
        """Return tasks whose dependencies are all completed."""
        return [
            task for task in self.tasks
            if not task.result
            and all(self.get_task(dep).result for dep in task.dependencies)
        ]

class MultiAgentOrchestrator:
    """Coordinates multiple agents working on a complex task."""

    def __init__(self, model: str = "gpt-5.6-sol"):
        self.model = model
        self.client = LLMClient()
        self.agents = {}
        self.results = {}

    def register_agent(self, agent: Agent):
        self.agents[agent.role] = agent

    async def plan_task(self, complex_task: str) -> TaskPlan:
        """Use the orchestrator agent to break down a task into sub-tasks."""
        planning_prompt = f"""Analyze this complex task and break it into sub-tasks.
For each sub-task, assign one of these roles: researcher, coder, reviewer, writer.
Define dependencies between tasks where appropriate.

Return a JSON array with this format:
[
  {{
    "id": "task_1",
    "description": "Clear description of what to do",
    "agent_role": "researcher",
    "dependencies": []
  }}
]

Task: {complex_task}
"""
        response = await self.client.complete(
            [{"role": "user", "content": planning_prompt}],
            model=self.model,
            temperature=0.3
        )

        # Parse the JSON response
        import re
        json_match = re.search(r'\[.*?\]', response, re.DOTALL)
        if json_match:
            tasks_data = json.loads(json_match.group())
            tasks = [
                Task(
                    id=t["id"],
                    description=t["description"],
                    agent_role=AgentRole(t["agent_role"]),
                    dependencies=t.get("dependencies", [])
                )
                for t in tasks_data
            ]
            return TaskPlan(tasks=tasks)
        raise ValueError("Failed to parse task plan from LLM response")

    async def execute_task(self, task: Task) -> str:
        """Execute a single task with the appropriate agent."""
        agent = self.agents.get(task.agent_role)
        if not agent:
            # Fall back to orchestrator
            agent = self.agents[AgentRole.ORCHESTRATOR]

        # Include context from dependency results
        context_parts = []
        for dep_id in task.dependencies:
            if dep_id in self.results:
                context_parts.append(
                    f"Results from {dep_id}:\n{self.results[dep_id]}"
                )

        full_task = task.description
        if context_parts:
            full_task = (
                "Context from previous tasks:\n\n"
                + "\n\n".join(context_parts)
                + f"\n\nYour task: {task.description}"
            )

        result = await agent.execute(full_task, stream=True)
        self.results[task.id] = result
        task.result = result
        return result

    async def run(self, complex_task: str) -> dict[str, str]:
        """Run the full multi-agent pipeline."""
        print(f"=== Multi-Agent System Starting ===")
        print(f"Task: {complex_task}\n")

        # Step 1: Plan
        print(">> Orchestrator is planning the task...")
        plan = await self.plan_task(complex_task)
        print(f">> Plan created with {len(plan.tasks)} tasks\n")

        # Step 2: Execute tasks (respecting dependencies)
        completed = set()
        while len(completed) < len(plan.tasks):
            ready = [
                t for t in plan.tasks
                if all(d in completed for d in t.dependencies)
                and t.id not in completed
            ]
            if not ready:
                break

            # Run ready tasks in parallel
            tasks_to_run = [self.execute_task(t) for t in ready]
            results = await asyncio.gather(*tasks_to_run)

            for i, result in enumerate(results):
                completed.add(ready[i].id)

        return self.results

async def main():
    # Create the orchestrator
    orchestrator = MultiAgentOrchestrator(model="gpt-5.6-sol")

    # Register specialist agents
    orchestrator.register_agent(Agent(
        name="Coordinator",
        role=AgentRole.ORCHESTRATOR,
        model="gpt-5.6-sol"
    ))
    orchestrator.register_agent(Agent(
        name="Research Bot",
        role=AgentRole.RESEARCHER,
        model="claude-sonnet-4-20250514"
    ))
    orchestrator.register_agent(Agent(
        name="Code Bot",
        role=AgentRole.CODER,
        model="gpt-5.6-sol"
    ))
    orchestrator.register_agent(Agent(
        name="Review Bot",
        role=AgentRole.REVIEWER,
        model="claude-opus-5"
    ))
    orchestrator.register_agent(Agent(
        name="Docs Bot",
        role=AgentRole.WRITER,
        model="deepseek-v4-pro"
    ))

    # Define the task
    task = """Build a Python REST API for a todo application with:
    - CRUD endpoints for tasks
    - JWT authentication
    - Input validation
    - Complete documentation"""

    # Run the multi-agent system
    results = await orchestrator.run(task)

    print("\n\n=== Final Results ===")
    for task_id, result in results.items():
        print(f"\n--- {task_id} ---")
        print(result[:500] + "..." if len(result) > 500 else result)

if __name__ == "__main__":
    asyncio.run(main())

async def execute_with_retry(agent, task, max_retries=3):
    for attempt in range(max_retries):
        try:
            return await agent.execute(task)
        except Exception as e:
            if attempt == max_retries - 1:
                raise
            await asyncio.sleep(2 ** attempt)

import json
from pathlib import Path

def save_checkpoint(orchestrator, path="checkpoints/"):
    Path(path).mkdir(exist_ok=True)
    for task_id, result in orchestrator.results.items():
        with open(f"{path}/{task_id}.json", "w") as f:
            json.dump({"task_id": task_id, "result": result}, f)

MAX_TOTAL_TOKENS = 500_000  # ~$1-2 depending on model
total_tokens = 0

async def bounded_execute(agent, task):
    global total_tokens
    if total_tokens >= MAX_TOTAL_TOKENS:
        raise RuntimeError("Token budget exceeded")
    result = await agent.execute(task)
    # Track tokens from response headers or usage data
    return result