Orchestrator-Worker

planner = Agent(
    name="ResearchPlanner",
    instructions="You are a research planner. Break the topic into 2-4 research sub-tasks.",
    output_type=TaskList,
)

researcher = Agent(
    name="Researcher",
    instructions="You are a research assistant. Provide a concise, factual answer.",
)

writer = Agent(
    name="ReportWriter",
    instructions="You are a report writer. Combine the research findings into a cohesive report.",
)

report_service = restate.Service("ResearchReport")


@report_service.handler()
async def generate(ctx: restate.Context, req: ReportRequest) -> dict:
    # Step 1: Orchestrator creates a research plan
    plan_result = await DurableRunner.run(planner, req.topic)
    tasks = plan_result.final_output.tasks

    # Step 2: Dispatch workers in parallel
    worker_promises = []
    for task in tasks:
        promise = ctx.service_call(run_researcher, task)
        worker_promises.append(promise)

    await restate.gather(*worker_promises)
    findings = [await p for p in worker_promises]

    # Step 3: Combine results into a report
    report_result = await DurableRunner.run(
        writer,
        f"Topic: {req.topic}\n\nResearch findings:\n{json.dumps(findings, indent=2)}",
    )

    return {"report": report_result.final_output, "task_count": len(tasks)}


researcher_service = restate.Service("Researcher")


@researcher_service.handler()
async def run_researcher(ctx: restate.Context, task: ResearchTask) -> str:
    result = await DurableRunner.run(researcher, task.question)
    return result.final_output

report_service = restate.VirtualObject("ResearchReport")


@report_service.handler()
async def generate(ctx: restate.ObjectContext, req: ReportRequest) -> dict:
    session_id = str(ctx.uuid())
    # Step 1: Orchestrator creates a research plan
    plan_events = plan_runner.run_async(
        user_id=ctx.key(),
        session_id=session_id,
        new_message=Content(role="user", parts=[Part.from_text(text=req.topic)]),
    )
    plan_output = await parse_agent_response(plan_events)
    tasks = TaskList.model_validate_json(plan_output).tasks

    # Step 2: Dispatch workers in parallel
    worker_promises = []
    for task in tasks:
        promise = ctx.object_call(run_researcher, key=str(ctx.uuid()), arg=task)
        worker_promises.append(promise)

    await restate.gather(*worker_promises)
    findings = [await p for p in worker_promises]

    # Step 3: Combine results into a report
    results = f"Topic: {req.topic}\n\nResearch findings:\n{json.dumps(findings)}"
    events = writer_runner.run_async(
        user_id=ctx.key(),
        session_id=session_id,
        new_message=Content(role="user", parts=[Part.from_text(text=results)]),
    )
    report = await parse_agent_response(events)

    return {"report": report, "task_count": len(tasks)}


researcher_service = restate.VirtualObject("Researcher")


@researcher_service.handler()
async def run_researcher(ctx: restate.ObjectContext, task: ResearchTask) -> str:
    events = research_runner.run_async(
        user_id=ctx.key(),
        session_id=str(ctx.uuid()),
        new_message=Content(role="user", parts=[Part.from_text(text=task.question)]),
    )
    return await parse_agent_response(events)

planner = Agent(
    "openai:gpt-5.4",
    system_prompt="You are a research planner. Break the topic into 2-4 research sub-tasks.",
    output_type=TaskList,
)
restate_planner = RestateAgent(planner)

researcher = Agent(
    "openai:gpt-5.4",
    system_prompt="You are a research assistant. Provide a concise, factual answer.",
)
restate_researcher = RestateAgent(researcher)

writer = Agent(
    "openai:gpt-5.4",
    system_prompt="You are a report writer. Combine the research findings into a cohesive report.",
)
restate_writer = RestateAgent(writer)

report_service = restate.Service("ResearchReport")


@report_service.handler()
async def generate(ctx: restate.Context, req: ReportRequest) -> dict:
    # Step 1: Orchestrator creates a research plan
    plan_result = await restate_planner.run(req.topic)
    tasks = plan_result.output.tasks

    # Step 2: Dispatch workers in parallel
    worker_promises = []
    for task in tasks:
        promise = ctx.service_call(run_researcher, task)
        worker_promises.append(promise)

    await restate.gather(*worker_promises)
    findings = [await p for p in worker_promises]

    # Step 3: Combine results into a report
    report_result = await restate_writer.run(
        f"Topic: {req.topic}\n\nResearch findings:\n{json.dumps(findings)}",
    )

    return {"report": report_result.output, "task_count": len(tasks)}


researcher_service = restate.Service("Researcher")


@researcher_service.handler()
async def run_researcher(_ctx: restate.Context, task: ResearchTask) -> str:
    result = await restate_researcher.run(task.question)
    return result.output

planner = create_agent(
    model=init_chat_model("openai:gpt-5.4"),
    system_prompt="You are a research planner. Break the topic into 2-4 research sub-tasks.",
    response_format=TaskList,
    middleware=[RestateMiddleware()],
)

researcher = create_agent(
    model=init_chat_model("openai:gpt-5.4"),
    system_prompt="You are a research assistant. Provide a concise, factual answer.",
    middleware=[RestateMiddleware()],
)

writer = create_agent(
    model=init_chat_model("openai:gpt-5.4"),
    system_prompt="You are a report writer. Combine the research findings into a cohesive report.",
    middleware=[RestateMiddleware()],
)

report_service = restate.Service("ResearchReport")


@report_service.handler()
async def generate(ctx: restate.Context, req: ReportRequest) -> dict:
    # Step 1: Orchestrator creates a research plan.
    plan_result = await planner.ainvoke({"messages": req.topic})
    tasks: list[ResearchTask] = plan_result["structured_response"].tasks

    # Step 2: Dispatch researchers in parallel.
    worker_promises = [ctx.service_call(run_researcher, task) for task in tasks]
    await restate.gather(*worker_promises)
    findings = [await p for p in worker_promises]

    # Step 3: Combine into a report.
    message = f"Topic: {req.topic}\n\nResearch findings:\n{json.dumps(findings)}"
    report_result = await writer.ainvoke({"messages": message})

    return {"report": report_result["messages"][-1].content, "task_count": len(tasks)}


researcher_service = restate.Service("Researcher")


@researcher_service.handler()
async def run_researcher(_ctx: restate.Context, task: ResearchTask) -> str:
    result = await researcher.ainvoke({"messages": task.question})
    return result["messages"][-1].content

export const researchWorker = restate.service({
  name: "ResearchWorker",
  handlers: {
    research: async (ctx: restate.Context, req: { question: string }) => {
      const answer = await ctx.run(
        "Research",
        async () =>
          llmCall(
            `You are a research assistant. Provide a concise, factual answer.\n\n${req.question}`,
          ),
        { maxRetryAttempts: 3 },
      );
      return { question: req.question, answer: answer.text };
    },
  },
});

const orchestrator = restate.service({
  name: "ResearchReport",
  handlers: {
    generate: restate.createServiceHandler(
        { input: schema(ResearchRequestSchema) },
        async (ctx: restate.Context, {topic}: { topic: string }) => {
      // Step 1: Orchestrator creates a research plan
      const planJson = await ctx.run(
        "Create research plan",
        async () =>
          llmCall(
            `You are a research planner. Break the topic into 2-4 research
          sub-tasks. Respond with a JSON array of strings, each a specific
          research question. Example: ["question 1", "question 2"]\n\nTopic: ${topic}`,
          ),
        { maxRetryAttempts: 3 },
      );
      const tasks: string[] = JSON.parse(planJson.text);

      // Step 2: Dispatch workers in parallel
      const workerResults = await RestatePromise.all(
        tasks.map((question) =>
          ctx.serviceClient(researchWorker).research({ question }),
        ),
      );

      // Step 3: Combine results into a report
      const report = await ctx.run(
        "Write report",
        async () =>
          llmCall(
            `You are a report writer. Combine the research findings into a cohesive report.\n\n
            Topic: ${topic}\n\nResearch findings:\n${JSON.stringify(workerResults)}`,
          ),
        { maxRetryAttempts: 3 },
      );

      return { report: report.text, taskCount: tasks.length };
        },
    ),
  },
});

researcher_service = restate.Service("ResearchWorker")


@researcher_service.handler()
async def research(ctx: restate.Context, req: ResearchTask) -> dict:
    answer = await ctx.run_typed(
        "Research",
        llm_call,
        RunOptions(max_attempts=3),
        messages=f"You are a research assistant. Provide a concise, factual answer. {req.question}",
    )
    return {"question": req.question, "answer": answer.content}


report_service = restate.Service("ResearchReport")


@report_service.handler()
async def generate(ctx: restate.Context, req: ReportRequest) -> dict:
    # Step 1: Orchestrator creates a research plan
    plan_result = await ctx.run_typed(
        "Create research plan",
        llm_call,
        RunOptions(max_attempts=3),
        messages=f"You are a research planner. Break the topic into 2-4 research sub-tasks. {req.topic}",
        response_format=TaskList,
    )
    if not plan_result.content:
        raise restate.TerminalError("No research plan created")
    tasks = TaskList.model_validate_json(plan_result.content).tasks

    # Step 2: Dispatch workers in parallel
    worker_promises = []
    for task in tasks:
        promise = ctx.service_call(research, task)
        worker_promises.append(promise)

    await restate.gather(*worker_promises)
    findings = [await p for p in worker_promises]

    # Step 3: Combine results into a report
    report = await ctx.run_typed(
        "Write report",
        llm_call,
        RunOptions(max_attempts=3),
        messages=f"You are a report writer. Combine the research findings into a cohesive report."
               f"Topic: {req.topic}\n\nResearch findings:\n{json.dumps(findings)}",
    )

    return {"report": report.content, "task_count": len(tasks)}