Spaces:
Runtime error
Runtime error
File size: 3,205 Bytes
283954c | 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 | import json
import time
from collections.abc import Iterator
from _core.llm import LLMClient
from _core.models import estimate_cost
from _core.tracer import Step
ORCH_PROMPT = (
"You are an orchestrator. Decompose the task into 2-4 independent subtasks, "
"each assigned to a named specialist. Output ONLY a JSON array of objects "
'{"role": str, "subtask": str}.'
)
SYNTH_PROMPT = (
"You are an orchestrator. Combine the workers' outputs into one coherent answer "
"to the original task."
)
class OrchestratorAgent:
def __init__(self, llm: LLMClient, max_workers: int = 4):
self.llm = llm
self.max_workers = max_workers
def _ask(self, system: str, user: str):
start = time.monotonic()
resp = self.llm.chat(
[{"role": "system", "content": system}, {"role": "user", "content": user}]
)
latency = int((time.monotonic() - start) * 1000)
cost = estimate_cost(self.llm.model, resp.prompt_tokens, resp.completion_tokens)
return resp, latency, cost
def run(self, task: str) -> Iterator[Step]:
resp, lat, cost = self._ask(ORCH_PROMPT, task)
try:
parsed = json.loads(resp.content or "[]")
except json.JSONDecodeError:
parsed = []
# Coerce whatever the model returned into a list of {role, subtask} dicts —
# small models often return a bare string array or a single object.
if not isinstance(parsed, list):
parsed = [parsed]
subtasks = []
for item in parsed[: self.max_workers]:
if isinstance(item, dict):
subtasks.append(
{"role": item.get("role", "specialist"), "subtask": item.get("subtask", task)}
)
else:
subtasks.append({"role": "generalist", "subtask": str(item)})
if not subtasks:
subtasks = [{"role": "generalist", "subtask": task}]
yield Step(
kind="thought",
content="Decomposition:\n"
+ "\n".join(f"- [{s['role']}] {s['subtask']}" for s in subtasks),
tokens=resp.prompt_tokens + resp.completion_tokens,
cost_usd=cost,
latency_ms=lat,
)
outputs: list[str] = []
for s in subtasks:
yield Step(kind="action", content=f"Dispatch to {s['role']}: {s['subtask']}")
wresp, wlat, wcost = self._ask(
f"You are a {s['role']}. Answer concisely.", s["subtask"]
)
yield Step(
kind="observation",
content=wresp.content or "",
tokens=wresp.prompt_tokens + wresp.completion_tokens,
cost_usd=wcost,
latency_ms=wlat,
)
outputs.append(f"[{s['role']}] {wresp.content}")
fresp, flat, fcost = self._ask(
SYNTH_PROMPT, f"Task: {task}\nWorker outputs:\n" + "\n".join(outputs)
)
yield Step(
kind="final",
content=fresp.content or "",
tokens=fresp.prompt_tokens + fresp.completion_tokens,
cost_usd=fcost,
latency_ms=flat,
)
|