Tour of Restate for Agents with OpenAI SDK

@function_tool(failure_error_function=raise_restate_errors)
async def get_weather(
    wrapper: RunContextWrapper[restate.Context], req: WeatherRequest
) -> WeatherResponse:
    """Get the current weather for a given city."""
    # Do durable steps using the Restate context
    restate_context = wrapper.context
    return await restate_context.run_typed("Get weather", fetch_weather, city=req.city)


weather_agent = Agent[restate.Context](
    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(restate_context: restate.Context, prompt: WeatherPrompt) -> str:

    result = await Runner.run(
        weather_agent,
        input=prompt.message,
        # Pass the Restate context to tools to make tool execution steps durable
        context=restate_context,
        # Choose any model and let Restate persist your calls
        run_config=RunConfig(
            model="gpt-4o",
            model_provider=DurableModelCalls(restate_context),
            model_settings=ModelSettings(parallel_tool_calls=False),
        ),
    )

    return result.final_output

@function_tool(failure_error_function=raise_restate_errors)
async def human_approval(
    wrapper: RunContextWrapper[restate.Context], claim: InsuranceClaim
) -> str:
    """Ask for human approval for high-value claims."""
    restate_context = wrapper.context

    # 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("requestApproval")
async def request_approval(
    restate_context: 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 = 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

@function_tool(failure_error_function=raise_restate_errors)
async def human_approval(
    wrapper: RunContextWrapper[restate.Context], claim: InsuranceClaim
) -> str:
    """Ask for human approval for high-value claims using sub-workflow."""
    restate_context = wrapper.context

    # Call the human approval sub-workflow
    return await restate_context.service_call(request_approval, claim)

chat = VirtualObject("Chat")


@chat.handler()
async def message(restate_context: ObjectContext, chat_message: ChatMessage) -> dict:

    restate_session = await RestateSession.create(
        session_id=restate_context.key(), ctx=restate_context
    )

    result = await Runner.run(
        Agent(name="Assistant", instructions="You are a helpful assistant."),
        input=chat_message.message,
        run_config=RunConfig(
            model="gpt-4o",
            model_provider=DurableModelCalls(restate_context),
            model_settings=ModelSettings(parallel_tool_calls=False),
        ),
        session=restate_session,
    )
    return result.final_output


@chat.handler(kind="shared")
async def get_history(ctx: ObjectSharedContext):
    return await ctx.get("items") or []

intake_agent = Agent[restate.ObjectContext](
    name="IntakeAgent",
    instructions="Route insurance claims to the appropriate specialist: medical, auto, or property.",
)

medical_specialist = Agent[restate.ObjectContext](
    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.",
)

auto_specialist = Agent[restate.ObjectContext](
    name="AutoSpecialist",
    handoff_description="I handle auto insurance claims from intake to final decision.",
    instructions="Assess auto claims for liability and damage. Approve/deny up to $25,000.",
)

# Configure handoffs so intake agent can route to specialists
intake_agent.handoffs = [medical_specialist, auto_specialist]

agent_dict = {
    "IntakeAgent": intake_agent,
    "MedicalSpecialist": medical_specialist,
    "AutoSpecialist": auto_specialist,
}

agent_service = restate.VirtualObject("MultiAgentClaimApproval")


@agent_service.handler()
async def run(restate_context: restate.ObjectContext, claim: InsuranceClaim) -> str:

    # Store context in Restate's key-value store
    last_agent_name = (
        await restate_context.get("last_agent_name", type_hint=str) or "IntakeAgent"
    )
    last_agent = agent_dict.get(last_agent_name, intake_agent)

    restate_session = await RestateSession.create(
        session_id=restate_context.key(), ctx=restate_context
    )
    result = await Runner.run(
        last_agent,
        input=f"Claim: {claim.model_dump_json()}",
        context=restate_context,
        run_config=RunConfig(
            model="gpt-4o",
            model_provider=DurableModelCalls(restate_context),
            model_settings=ModelSettings(parallel_tool_calls=False),
        ),
        session=restate_session,
    )

    restate_context.set("last_agent_name", result.last_agent.name)

    return result.final_output

# Durable service call to the fraud agent; persisted and retried by Restate
@function_tool(failure_error_function=raise_restate_errors)
async def check_fraud(
    wrapper: RunContextWrapper[restate.Context], claim: InsuranceClaim
) -> str:
    """Analyze the probability of fraud."""
    restate_context = wrapper.context
    return await restate_context.service_call(run_fraud_agent, claim)


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

agent_service = restate.Service("RemoteMultiAgentClaimApproval")


@agent_service.handler()
async def run(restate_context: restate.Context, claim: InsuranceClaim) -> str:
    result = await Runner.run(
        claim_approval_coordinator,
        input=f"Claim: {claim.model_dump_json()}",
        context=restate_context,
        run_config=RunConfig(
            model="gpt-4o",
            model_provider=DurableModelCalls(restate_context),
            model_settings=ModelSettings(parallel_tool_calls=False),
        ),
    )
    return result.final_output

@function_tool(failure_error_function=raise_restate_errors)
async def calculate_metrics(
    wrapper: RunContextWrapper[restate.Context], claim: InsuranceClaim
) -> list[str]:
    """Calculate claim metrics."""
    restate_context = wrapper.context

    # Run tools/steps in parallel with durable execution
    results_done = await restate.gather(
        restate_context.run_typed("eligibility", check_eligibility, claim=claim),
        restate_context.run_typed("cost", compare_to_standard_rates, claim=claim),
        restate_context.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 Runner.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}",
        run_config=RunConfig(
            model="gpt-4o",
            model_provider=DurableModelCalls(restate_context),
            model_settings=ModelSettings(parallel_tool_calls=False),
        ),
    )
    return result.final_output

@function_tool(failure_error_function=raise_restate_errors)
async def get_weather(
    wrapper: RunContextWrapper[restate.Context], req: WeatherRequest
) -> WeatherResponse:
    """Get the current weather for a given city."""
    restate_context = wrapper.context
    return await restate_context.run_typed("Get weather", fetch_weather, city=req.city)

try:
    result = await Runner.run(
        weather_agent,
        input=prompt.message,
        context=restate_context,
        run_config=RunConfig(
            model="gpt-4o",
            model_provider=DurableModelCalls(restate_context, max_retries=2),
            model_settings=ModelSettings(parallel_tool_calls=False),
        ),
    )
except restate.TerminalError as e:
    # Handle terminal errors gracefully
    return "The agent couldn't complete the request."