Tour of Restate for Agents with OpenAI SDK

@durable_function_tool
async def get_weather(city: WeatherRequest) -> WeatherResponse:
    """Get the current weather for a given city."""
    return await restate_context().run_typed("get weather", fetch_weather, req=city)


agent = Agent(
    name="WeatherAgent",
    instructions="You are a helpful agent that provides weather updates.",
    tools=[get_weather],
)

agent_service = restate.Service("WeatherAgent")


@agent_service.handler()
async def run(_ctx: restate.Context, req: WeatherPrompt) -> str:
    result = await DurableRunner.run(agent, req.message)
    return result.final_output

@durable_function_tool
async def human_approval(claim: InsuranceClaim) -> str:
    """Ask for human approval for high-value claims."""

    # Create an awakeable for human approval
    approval_id, approval_promise = restate_context().awakeable(type_hint=str)

    # Request human review
    await restate_context().run_typed(
        "Request review", request_human_review, claim=claim, awakeable_id=approval_id
    )

    # Wait for human approval
    return await approval_promise

# Sub-workflow service for human approval
human_approval_workflow = restate.Service("HumanApprovalWorkflow")


@human_approval_workflow.handler()
async def review(ctx: restate.Context, claim: InsuranceClaim) -> str:
    """Request human approval for a claim and wait for response."""
    # Create an awakeable that can be resolved via HTTP
    approval_id, approval_promise = ctx.awakeable(type_hint=str)

    # Request human review
    await ctx.run_typed(
        "Request review", request_human_review, claim=claim, awakeable_id=approval_id
    )

    # Wait for human approval
    return await approval_promise

@durable_function_tool
async def human_approval(claim: InsuranceClaim) -> str:
    """Ask for human approval for high-value claims."""
    return await restate_context().service_call(review, claim)

chat = VirtualObject("Chat")


@chat.handler()
async def message(_ctx: ObjectContext, req: ChatMessage) -> dict:
    # Set use_restate_session=True to store the session in Restate's key-value store
    # Make sure you use a VirtualObject to enable this feature
    result = await DurableRunner.run(
        Agent(name="Assistant", instructions="You are a helpful assistant."),
        req.message,
        use_restate_session=True,
    )
    return result.final_output

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,
    "AutoSpecialist": 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

# Durable service call to the fraud agent; persisted and retried by Restate
@durable_function_tool
async def check_fraud(claim: InsuranceClaim) -> str:
    """Analyze the probability of fraud."""
    return await restate_context().service_call(run_fraud_agent, claim)


agent = Agent(
    name="ClaimApprovalCoordinator",
    instructions="You are a claim approval engine. Analyze the claim and use your tools to decide whether to approve it.",
    tools=[check_eligibility, check_fraud],
)

agent_service = restate.Service("RemoteMultiAgentClaimApproval")


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

@durable_function_tool
async def calculate_metrics(claim: InsuranceClaim) -> list[str]:
    """Calculate claim metrics."""
    ctx = restate_context()

    # Run tools/steps in parallel with durable execution
    results_done = await restate.gather(
        ctx.run_typed("eligibility", check_eligibility, claim=claim),
        ctx.run_typed("cost", compare_to_standard_rates, claim=claim),
        ctx.run_typed("fraud", check_fraud, claim=claim),
    )
    return [await result for result in results_done]

@agent_service.handler()
async def run(restate_context: restate.Context, claim: InsuranceClaim) -> str:
    # Start multiple agents in parallel with auto retries and recovery
    eligibility = restate_context.service_call(run_eligibility_agent, claim)
    cost = restate_context.service_call(run_rate_comparison_agent, claim)
    fraud = restate_context.service_call(run_fraud_agent, claim)

    # Wait for all responses
    await restate.gather(eligibility, cost, fraud)

    # Run decision agent on outputs
    result = await DurableRunner.run(
        Agent(
            name="ClaimApprovalAgent", instructions="You are a claim decision engine."
        ),
        input=f"Decide about claim: {claim.model_dump_json()}. "
        "Base your decision on the following analyses:"
        f"Eligibility: {await eligibility} Cost {await cost} Fraud: {await fraud}",
    )
    return result.final_output

try:
    result = await DurableRunner.run(
        agent,
        req.message,
        llm_retry_opts=LlmRetryOpts(
            max_attempts=3, initial_retry_interval=timedelta(seconds=2)
        ),
    )
except restate.TerminalError as e:
    # Handle terminal errors gracefully
    return f"The agent couldn't complete the request: {e.message}"