Supervisor and workers: route subtasks to specialists

"""
Supervisor and workers: route subtasks to specialists

A supervisor decomposes a request into subtasks, assigns each to a named
specialist, and a synthesiser merges the results. Each specialist runs with
its own instruction and its own slice of context.

    Request -> Plan (typed subtasks) -> Specialists -> Synthesis

The real reason to reach for this is context isolation, not intelligence:
each worker sees only what its subtask needs, so one long transcript does not
drag every step. It is not a general upgrade. Most requests are better served
by one loop, and a team of agents that does not need to be a team is just a
larger bill and more places to fail.

Reach for it when the subtasks are genuinely independent, or when they need
materially different instructions. Not when they merely sound different.

Run: python supervisor_workers.py
"""

from typing import Literal

from dotenv import load_dotenv
from openai import OpenAI
from pydantic import BaseModel, Field

load_dotenv()

client = OpenAI()

MODEL = "gpt-5.4-mini"
MAX_SUBTASKS = 4

#: Each specialist is an instruction, not a model. Adding one is a line here.
SPECIALISTS: dict[str, str] = {
    "technical": (
        "You assess technical feasibility and implementation risk. Name the "
        "specific mechanism that would break, not general caution."
    ),
    "financial": (
        "You assess cost and commercial impact. Give ranges and state what "
        "drives them. Never invent a figure that was not provided."
    ),
    "operational": (
        "You assess who has to run this once it exists, and what they need. "
        "Focus on the ongoing burden rather than the build."
    ),
}


class Subtask(BaseModel):
    specialist: Literal["technical", "financial", "operational"]
    question: str = Field(
        description="The single question this specialist should answer."
    )


class Plan(BaseModel):
    subtasks: list[Subtask] = Field(
        description=(
            "One subtask per specialist that is genuinely needed. Do not "
            "include a specialist whose lens adds nothing to this request."
        )
    )
    reasoning: str = Field(description="Why these specialists and not the others.")


def plan(request: str) -> Plan:
    completion = client.beta.chat.completions.parse(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Decompose the request into subtasks for the specialists "
                    "listed. Use only the specialists that genuinely add "
                    "something. Fewer is better: a specialist with nothing "
                    "useful to say costs a call and dilutes the synthesis.\n\n"
                    "Specialists:\n"
                    + "\n".join(f"- {name}: {role}" for name, role in SPECIALISTS.items())
                ),
            },
            {"role": "user", "content": request},
        ],
        response_format=Plan,
    )
    return completion.choices[0].message.parsed


def run_specialist(name: str, question: str) -> str:
    """One worker. Sees its own question, not the whole transcript."""
    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": SPECIALISTS[name]},
            {"role": "user", "content": question},
        ],
    )
    return completion.choices[0].message.content or ""


def synthesise(request: str, findings: list[tuple[str, str]]) -> str:
    body = "\n\n".join(f"[{name}]\n{text}" for name, text in findings)
    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Merge the specialist findings into one answer. Attribute "
                    "each point to the specialist it came from. Where two "
                    "specialists disagree, say so rather than averaging them."
                ),
            },
            {"role": "user", "content": f"Request: {request}\n\n{body}"},
        ],
    )
    return completion.choices[0].message.content or ""


def run(request: str) -> str:
    print("  planning...")
    proposal = plan(request)
    print(f"      rationale: {proposal.reasoning}")

    subtasks = proposal.subtasks[:MAX_SUBTASKS]
    if not subtasks:
        # A plan with no subtasks is a real answer: this request did not need
        # a team. Say so rather than inventing work for the specialists.
        message = "No specialist lens applied to this request. Answer it directly."
        print(f"      {message}")
        return message

    findings: list[tuple[str, str]] = []
    for index, subtask in enumerate(subtasks, start=1):
        # Printed before the call, so a slow specialist reads as working.
        print(f"  [{index}/{len(subtasks)}] {subtask.specialist}: {subtask.question}")
        answer = run_specialist(subtask.specialist, subtask.question)
        print(f"      -> {answer[:80]}...")
        findings.append((subtask.specialist, answer))

    print("  synthesising...")
    result = synthesise(request, findings)
    print(f"      {result}")
    return result


if __name__ == "__main__":
    request = (
        "We are considering moving our nightly batch reconciliation to a "
        "streaming pipeline. Should we?"
    )
    print(f"\n> {request}")
    run(request)