AI Agent Orchestration: Designing and Implementing Multi-Agent Systems

import json
import re
 
def extract_json(text: str) -> dict:
    """Pull the first JSON object/array out of an LLM response."""
    fence = chr(96) * 3  # a code fence: three backticks
    body = text.strip()
 
    fenced = re.search(fence + r"(?:json)?\s*(.+?)\s*" + fence, body, re.S)
    if fenced:
        body = fenced.group(1).strip()
 
    starts = [i for i in (body.find("{"), body.find("[")) if i != -1]
    if not starts:
        raise ValueError(f"JSON not found in model output: {text[:200]}")
 
    # raw_decode tolerates prose trailing after the JSON
    obj, _ = json.JSONDecoder().raw_decode(body[min(starts):])
    return obj

from anthropic import Anthropic
import json
 
client = Anthropic()
 
def decompose_task(task: str, available_agents: list[dict]) -> list[dict]:
    """
    Use an LLM to decompose a task and assign sub-tasks to agents
    """
    decompose_prompt = f"""
    You are a task decomposition expert. Break the following task into sub-tasks
    that can be executed by the available agents.
    
    Main task: {task}
    
    Available agents:
    {json.dumps(available_agents, indent=2)}
    
    Respond in this JSON format:
    {{
        "subtasks": [
            {{
                "id": "task_1",
                "description": "Sub-task description",
                "assigned_agent": "agent_id",
                "depends_on": [],
                "can_parallel": true
            }}
        ],
        "execution_order": [["task_1", "task_2"], ["task_3"]]
    }}
    
    execution_order is a 2D array where each inner array contains tasks that can run in parallel.
    """
    
    response = client.messages.create(
        model="claude-opus-5",
        max_tokens=4096,
        messages=[{"role": "user", "content": decompose_prompt}]
    )
    
    return extract_json(response.content[0].text)
 
import concurrent.futures
 
def execute_task_plan(plan: dict, agents: dict) -> dict:
    """Execute a task plan sequentially and in parallel"""
    results: dict = {}
    failures: dict = {}
    
    for parallel_group in plan["execution_order"]:
        if len(parallel_group) == 1:
            task_id = parallel_group[0]
            subtask = next(t for t in plan["subtasks"] if t["id"] == task_id)
            agent = agents[subtask["assigned_agent"]]
            context = {dep: results[dep] for dep in subtask["depends_on"]}
            results[task_id] = agent.execute(subtask["description"], context)
        else:
            # Parallel execution — one failure must not erase the others
            with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
                futures = {}
                for task_id in parallel_group:
                    subtask = next(t for t in plan["subtasks"] if t["id"] == task_id)
                    agent = agents[subtask["assigned_agent"]]
                    context = {dep: results[dep] for dep in subtask["depends_on"]}
                    futures[executor.submit(
                        agent.execute, subtask["description"], context
                    )] = task_id
 
                for future in concurrent.futures.as_completed(futures):
                    task_id = futures[future]
                    try:
                        results[task_id] = future.result()
                    except Exception as exc:
                        failures[task_id] = f"{type(exc).__name__}: {exc}"
    
    return {"results": results, "failures": failures}

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict
import threading
 
@dataclass
class SharedWorkspace:
    """Workspace shared across all agents"""
    artifacts: Dict[str, Any] = field(default_factory=dict)
    messages: list = field(default_factory=list)
    _lock: threading.RLock = field(default_factory=threading.RLock)
    
    def write_artifact(self, key: str, value: Any, agent_id: str):
        """Thread-safe artifact write"""
        with self._lock:
            self.artifacts[key] = {
                "value": value,
                "written_by": agent_id,
                "timestamp": datetime.now(timezone.utc).isoformat()
            }
    
    def read_artifact(self, key: str) -> Any:
        with self._lock:
            item = self.artifacts.get(key)
            return item["value"] if item else None
    
    def post_message(self, from_agent: str, to_agent: str, content: str):
        with self._lock:
            self.messages.append({
                "from": from_agent,
                "to": to_agent,
                "content": content,
                "timestamp": datetime.now(timezone.utc).isoformat()
            })

def quality_gate(output: str, criteria: dict) -> dict:
    """Evaluate agent output against quality criteria"""
    evaluation_prompt = f"""
    Evaluate the following output against the given criteria.
    
    Output:
    {output}
    
    Criteria:
    {json.dumps(criteria)}
    
    Respond in JSON:
    {{
        "passed": true/false,
        "score": 0-100,
        "issues": ["issue 1", "issue 2"],
        "improvement_suggestions": ["suggestion 1", "suggestion 2"]
    }}
    """
    
    response = client.messages.create(
        model="claude-opus-5",
        max_tokens=1024,
        messages=[{"role": "user", "content": evaluation_prompt}]
    )
    
    return extract_json(response.content[0].text)
 
def iterative_refinement(task: str, agent, criteria: dict, max_iterations: int = 3) -> str:
    """Iteratively refine output until quality criteria are met"""
    result = agent.execute(task)
    
    for i in range(max_iterations):
        quality = quality_gate(result, criteria)
        if quality["passed"]:
            return result
        
        improvement_task = f"""
        Original task: {task}
        
        Previous output:
        {result}
        
        Areas to improve:
        {json.dumps(quality["improvement_suggestions"])}
        
        Generate an improved output addressing these points.
        """
        result = agent.execute(improvement_task)
    
    return result

class BudgetAwareAgent:
    def __init__(self, max_tokens: int = 4096, model: str = "claude-sonnet-5"):
        self.max_tokens = max_tokens
        self.model = model
        self.tokens_used = 0
    
    def execute(self, task: str) -> str:
        if self.tokens_used > self.max_tokens * 0.9:
            self.model = "claude-haiku-4-5"  # Fall back to lighter model
        
        response = client.messages.create(
            model=self.model,
            max_tokens=min(2048, self.max_tokens - self.tokens_used),
            messages=[{"role": "user", "content": task}]
        )
        
        self.tokens_used += response.usage.input_tokens + response.usage.output_tokens
        return response.content[0].text

import asyncio
 
async def execute_with_timeout(agent, task: str, timeout: float = 30.0) -> str:
    try:
        return await asyncio.wait_for(
            asyncio.to_thread(agent.execute, task),
            timeout=timeout
        )
    except asyncio.TimeoutError:
        return f"ERROR: Agent timed out after {timeout}s"