Multi-Agent Orchestration

medical_agent = Agent(
    name="MedicalSpecialist",
    handoff_description="I handle medical insurance claims from intake to final decision.",
    instructions="Review medical claims for coverage and necessity. Approve/deny up to $50,000.",
)

car_agent = Agent(
    name="CarSpecialist",
    handoff_description="I handle car insurance claims from intake to final decision.",
    instructions="Assess car claims for liability and damage. Approve/deny up to $25,000.",
)


intake_agent = Agent(
    name="IntakeAgent",
    instructions="Route insurance claims to the appropriate specialist",
    handoffs=[medical_agent, car_agent],
)

agent_dict = {
    "IntakeAgent": intake_agent,
    "MedicalSpecialist": medical_agent,
    "CarSpecialist": car_agent,
}

agent_service = restate.VirtualObject("MultiAgentClaimApproval")


@agent_service.handler()
async def run(ctx: restate.ObjectContext, claim: InsuranceClaim) -> str:
    # Store context in Restate's key-value store
    last_agent_name = await ctx.get("last_agent_name", type_hint=str) or "IntakeAgent"
    last_agent = agent_dict.get(last_agent_name, intake_agent)

    result = await DurableRunner.run(
        last_agent, f"Claim: {claim.model_dump_json()}", session=RestateSession()
    )

    ctx.set("last_agent_name", result.last_agent.name)
    return result.final_output

# AGENTS
# Determine which specialist to use based on claim type
medical_agent = Agent(
    model="gemini-2.5-flash",
    name="medical_specialist",
    description="Reviews medical insurance claims for coverage and necessity.",
    instruction="Review medical claims for coverage and necessity. Approve/deny up to $50,000.",
)

car_agent = Agent(
    model="gemini-2.5-flash",
    name="car_specialist",
    description="Assesses car insurance claims for liability and damage.",
    instruction="Assess car claims for liability and damage. Approve/deny up to $25,000.",
)

agent = Agent(
    model="gemini-2.5-flash",
    name="intake_agent",
    instruction="Route insurance claims to the appropriate specialist",
    sub_agents=[car_agent, medical_agent],
)

# Enables retries and recovery for model calls and tool executions
app = App(name=APP_NAME, root_agent=agent, plugins=[RestatePlugin()])
runner = Runner(app=app, session_service=RestateSessionService())

agent_service = restate.VirtualObject("MultiAgentClaimApproval")


@agent_service.handler()
async def run(ctx: restate.ObjectContext, claim: InsuranceClaim) -> str | None:
    events = runner.run_async(
        user_id=ctx.key(),
        session_id=claim.session_id,
        new_message=Content(
            role="user",
            parts=[Part.from_text(text=f"Claim: {claim.model_dump_json()}")],
        ),
    )
    return await parse_agent_response(events)

medical_agent = Agent(
    "openai:gpt-5.4",
    system_prompt="Review medical claims for coverage and necessity. Approve/deny up to $50,000.",
)
restate_medical_agent = RestateAgent(medical_agent)

car_agent = Agent(
    "openai:gpt-5.4",
    system_prompt="Assess car claims for liability and damage. Approve/deny up to $25,000.",
)
restate_car_agent = RestateAgent(car_agent)

intake_agent = Agent(
    "openai:gpt-5.4",
    system_prompt="Route insurance claims to the appropriate specialist using the available tools.",
)


@intake_agent.tool
async def consult_medical_specialist(
    _run_ctx: RunContext[None], claim: InsuranceClaim
) -> str:
    """Route to the medical specialist for medical insurance claims."""
    result = await restate_medical_agent.run(claim.model_dump_json())
    return result.output


@intake_agent.tool
async def consult_car_specialist(
    _run_ctx: RunContext[None], claim: InsuranceClaim
) -> str:
    """Route to the car specialist for car insurance claims."""
    result = await restate_car_agent.run(claim.model_dump_json())
    return result.output


restate_intake_agent = RestateAgent(intake_agent)
agent_service = restate.Service("MultiAgentClaimApproval")


@agent_service.handler()
async def run(_ctx: restate.ObjectContext, claim: InsuranceClaim) -> str:
    result = await restate_intake_agent.run(f"Claim: {claim.model_dump_json()}")
    return result.output

medical_agent = create_agent(
    model=init_chat_model("openai:gpt-5.4"),
    system_prompt=(
        "You are a medical insurance specialist. Review medical claims for "
        "coverage and necessity. Approve/deny up to $50,000."
    ),
    middleware=[RestateMiddleware()],
)

car_agent = create_agent(
    model=init_chat_model("openai:gpt-5.4"),
    system_prompt=(
        "You are a car insurance specialist. Assess car claims for liability "
        "and damage. Approve/deny up to $25,000."
    ),
    middleware=[RestateMiddleware()],
)


@tool
async def to_medical_specialist(claim_json: str) -> str:
    """Hand the claim to the medical specialist for evaluation."""
    result = await medical_agent.ainvoke({"messages": claim_json})
    return result["messages"][-1].content


@tool
async def to_car_specialist(claim_json: str) -> str:
    """Hand the claim to the car specialist for evaluation."""
    result = await car_agent.ainvoke({"messages": claim_json})
    return result["messages"][-1].content


intake_agent = create_agent(
    model=init_chat_model("openai:gpt-5.4"),
    tools=[to_medical_specialist, to_car_specialist],
    system_prompt=(
        "You are an intake agent. Route insurance claims to the appropriate "
        "specialist. Always call exactly one specialist tool, then summarize "
        "their decision."
    ),
    middleware=[RestateMiddleware()],
)


agent_service = restate.VirtualObject("MultiAgentClaimApproval")


@agent_service.handler()
async def run(ctx: restate.ObjectContext, claim: InsuranceClaim) -> str:
    history = await ctx.get("messages", type_hint=ChatHistory) or ChatHistory()
    history.messages.append(HumanMessage(content=f"Claim: {claim.model_dump_json()}"))

    result = await intake_agent.ainvoke({"messages": history.messages})

    ctx.set("messages", ChatHistory(messages=result["messages"]))
    return result["messages"][-1].content

const SPECIALISTS = {
  billingAgent: {
    description: "Expert in payments, charges, and refunds",
    prompt:
      "You are a billing support agent specializing in payments, charges, and refunds.",
  },
  accountAgent: {
    description: "Expert in login issues and security",
    prompt:
      "You are an account support agent specializing in login issues and security.",
  },
  productAgent: {
    description: "Expert in features and how-to guides",
    prompt:
      "You are a product support agent specializing in features and how-to guides.",
  },
} as const;

type Specialist = keyof typeof SPECIALISTS;

async function answer(ctx: Context, { message }: { message: string }) {
  // 1. First, decide if a specialist is needed
  const routingDecision = await ctx.run(
    "Pick specialist",
    // Use your preferred LLM SDK here - specify agents as tools
    async () => llmCall(message, createTools(SPECIALISTS)),
    { maxRetryAttempts: 3 },
  );

  // 2. No specialist needed? Give a general answer
  if (!routingDecision.toolCalls || routingDecision.toolCalls.length === 0) {
    return routingDecision.text;
  }

  // 3. Get the specialist's name
  const specialist = routingDecision.toolCalls[0].toolName as Specialist;

  // 4. Ask the specialist to answer
  const { text } = await ctx.run(
    `Ask ${specialist}`,
    async () =>
      llmCall([
        { role: "user", content: message },
        { role: "system", content: SPECIALISTS[specialist].prompt },
      ]),
    { maxRetryAttempts: 3 },
  );

  return text;
}

# Create the routing service
router = restate.Service("AgentRouter")

# Our team of AI specialists
SPECIALISTS = {
    "BillingAgent": "Expert in payments, charges, and refunds",
    "AccountAgent": "Expert in login issues and security",
    "ProductAgent": "Expert in features and how-to guides",
}


@router.handler()
async def answer(ctx: restate.Context, question: Question) -> str | None:
    """Classify request and route to appropriate specialized agent."""

    # 1. First, decide if a specialist is needed
    routing_decision = await ctx.run_typed(
        "Pick specialist",
        llm_call,  # Use your preferred LLM SDK here
        RunOptions(max_attempts=3),
        messages=f"""You are a customer service routing system. 
        Choose the appropriate specialist, or respond directly if no specialist is needed. 
        {question.message}""",
        tools=[tool(name=name, description=desc) for name, desc in SPECIALISTS.items()],
    )

    # 2. No specialist needed? Give a general answer
    if not routing_decision.tool_calls:
        return routing_decision.content

    # 3. Get the specialist's name
    specialist = routing_decision.tool_calls[0].function.name or "ProductAgent"

    # 4. Ask the specialist to answer
    response = await ctx.run_typed(
        f"Ask {specialist}",
        llm_call,
        RunOptions(max_attempts=3),
        messages=f"""You are a {SPECIALISTS.get(specialist)} specialist."
        Answer the question: {question.message}""",
    )

    return response.content