Spaces:
Sleeping
Sleeping
File size: 5,527 Bytes
936905f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """
orchestrator.py — Task decomposition, parallel dispatch, and synthesis.
Measures real parallel vs sequential timing and returns the speedup ratio.
"""
import asyncio
import json
import logging
import time
import uuid
from openai import AsyncOpenAI
from agents import AGENT_MAP
from config import get_settings
logger = logging.getLogger(__name__)
settings = get_settings()
client = AsyncOpenAI(api_key=settings.openai_api_key)
DECOMPOSE_SYSTEM = """You are a task decomposition engine for a multi-agent AI system.
Break the user's task into 3-5 focused subtasks. Assign each to the best agent type.
Agent types and their strengths:
- research: facts, comparisons, background knowledge
- code: implementation, architecture, code review, technical specs
- analysis: tradeoffs, evaluation, structured reasoning
- writer: documentation, summaries, checklists, final prose
Return ONLY a valid JSON array. No markdown fences. No explanation.
Each item must have exactly these keys:
"id": integer starting at 1
"agent_type": one of: research, code, analysis, writer
"prompt": the specific subtask instruction (be precise and self-contained)
"depends_on": [] (empty list — all run in parallel for now)
Example output:
[
{"id": 1, "agent_type": "research", "prompt": "Research X", "depends_on": []},
{"id": 2, "agent_type": "code", "prompt": "Write code for Y", "depends_on": []}
]"""
async def decompose_task(task: str) -> list[dict]:
"""Call OpenAI to break the task into structured subtasks."""
response = await client.chat.completions.create(
model=settings.model_name,
messages=[
{"role": "system", "content": DECOMPOSE_SYSTEM},
{"role": "user", "content": task},
],
temperature=0.2, # Low temp for consistent structure
)
raw = response.choices[0].message.content.strip()
# Strip markdown fences if the model ignores the instruction
if raw.startswith("```"):
raw = raw.split("```")[1]
if raw.startswith("json"):
raw = raw[4:]
raw = raw.strip()
subtasks = json.loads(raw)
logger.info("Decomposed into %d subtasks", len(subtasks))
return subtasks
async def run_parallel(subtasks: list[dict], original_task: str) -> dict:
"""
Run all subtasks in parallel using asyncio.gather().
Measures both parallel wall time and estimated sequential time.
Returns results keyed by subtask id.
"""
async def run_one(subtask: dict) -> tuple[str, dict]:
agent_fn = AGENT_MAP.get(subtask["agent_type"])
if agent_fn is None:
logger.warning("Unknown agent type: %s — falling back to research", subtask["agent_type"])
agent_fn = AGENT_MAP["research"]
result = await agent_fn(subtask["prompt"], original_task)
return str(subtask["id"]), result
parallel_start = time.perf_counter()
pairs = await asyncio.gather(*[run_one(s) for s in subtasks])
parallel_time = time.perf_counter() - parallel_start
results = {task_id: data for task_id, data in pairs}
# Sequential time = sum of all individual times (what it would have taken one-by-one)
sequential_time = sum(r["time_taken"] for r in results.values())
speedup = round(sequential_time / parallel_time, 2) if parallel_time > 0 else 1.0
return {
"results": results,
"parallel_time": round(parallel_time, 3),
"sequential_time": round(sequential_time, 3),
"speedup": speedup,
}
async def synthesize(task: str, subtask_results: dict) -> str:
"""Combine all agent results into a final cohesive answer."""
combined = "\n\n".join(
f"--- {data['agent'].upper()} AGENT ---\n{data['result']}"
for data in sorted(subtask_results.values(), key=lambda x: x.get("id", 0))
)
response = await client.chat.completions.create(
model=settings.model_name,
messages=[
{
"role": "system",
"content": (
"You are a synthesis agent. Integrate outputs from multiple specialist agents "
"into a single cohesive, well-structured final answer. "
"Do not concatenate — weave insights together. Use clear headers."
),
},
{
"role": "user",
"content": (
f"Original task: {task}\n\n"
f"Agent outputs:\n{combined}\n\n"
"Write the final synthesized answer."
),
},
],
temperature=0.4,
)
return response.choices[0].message.content.strip()
async def run_pipeline(task: str) -> dict:
"""
Full pipeline: decompose → parallel agents → synthesize.
Returns everything the API endpoint needs.
"""
session_id = str(uuid.uuid4())[:8]
subtasks = await decompose_task(task)
plan = "\n".join(
f" [{s['agent_type'].upper()}] {s['prompt']}" for s in subtasks
)
dispatch = await run_parallel(subtasks, task)
subtask_results = dispatch["results"]
final_answer = await synthesize(task, subtask_results)
return {
"session_id": session_id,
"plan": plan,
"subtask_results": subtask_results,
"final_answer": final_answer,
"parallel_time": dispatch["parallel_time"],
"sequential_time": dispatch["sequential_time"],
"speedup": dispatch["speedup"],
"agents_used": len(subtasks),
}
|