[RETIRED] Restate AI Agents examples

This repo is outdated. Look at this one instead: https://github.com/restatedev/ai-examples

[RETIRED] Restate AI Agents examples

This repository contains examples of how to use Restate for LLM-based / AI Agent use cases.

Restate gives many benefits when implementing LLM-based applications and agents. In this ReadMe we describe a few options on how to integrate Restate with existing AI SDKs.

The integration goes from loose integration (calling LLMs or Agent SDKs from Restate actions) to tight integration (implementing long-running multi-agent setups with Restate constructs).

At this point, we see each of these proposed integrations as a valid way of using Restate for AI applications, depending on the use case and desired behavior.

We are interested in understanding which of these appeals most to users.

The options:

Using Restate for orchestrating LLM calls
Using Restate for stateful, explicit control flow over LLM feedback loops
Using Restate together with Agent SDKs for scalable, resilient agent loops
Using Restate for agentic workflows
Mixing static, code-defined workflows with agentic workflows
Long-lived multi-agent setups

Using Restate for orchestrating LLM calls

This example shows how to implement Anthropic's agents cookbook patterns using Restate.

The patterns in this example show how you can use Restate to harden LLM-based routing decisions and tool executions with Restate.

The patterns included here:

For example, parallelization can be implemented as follows:

parallelization_svc = restate.Service("ParallelizationService")


class ParallelizationRequest(BaseModel):
    prompt: str
    inputs: list[str]


@parallelization_svc.handler()
async def run_in_parallel(
    ctx: restate.Context, req: ParallelizationRequest
) -> list[str]:
    futures = [
        ctx.run(
            f"LLM call {item}",
            lambda item=item: llm_call(f"{req.prompt}\nInput: {item}"),
        )
        for item in req.inputs
    ]
    results_done = await restate.gather(*futures)
    results = [await result for result in results_done]
    return results

The benefits of using Restate here are:

Automatic retries of failed tasks: LLM API down, timeouts, long-running tasks, infrastructure failures, etc.
Recovery of previous progress: After a failure, Restate recovers the progress the execution did before the cash. It persists routing decisions, tool execution outcomes, and deterministically replays them after failures. As opposed to executing all tasks again, Restate only re-executes the tasks that were not completed before the failure.
Management of the orchestration task and all parallel subtasks: Restate guarantees all tasks run to completion exactly once.
Idempotency: Restate can deduplicate requests and tool executions.

Note that this does not use an Agent SDK. These patterns only use an SDK to help with calling the LLM: OpenAI Responses API, or Anthropic's API.

Run the examples

Using Restate for stateful, explicit control flow and LLM feedback loops

You can integrate Restate one step deeper into your LLM-based apps by using Restate's stateful entities, called Virtual Objects, to manage the session state of the interactions with the LLM.

In the previous section, we talked about the evaluator-optimizer pattern which implements a loop of letting the LLM generate a response, and then asking the LLM to evaluate the response, then generate a new response based on the evaluation, and so on. This is a feedback loop that can be used to improve the quality of the LLM's output.

Instead of letting another LLM do the evaluation, you can also ask a human to evaluate the output of the LLM.

In this example, we track the generated responses and the chain of thought in a Restate Virtual Object, and let the user repeatedly invoke the handler with new insights (full code example)

human_evaluator_optimizer = restate.VirtualObject("HumanEvaluatorOptimizer")


@human_evaluator_optimizer.handler()
async def run(ctx: restate.ObjectContext, req: MyRequest) -> tuple[str, list[dict]]:
    memory = await ctx.get("memory") or []
    chain_of_thought = await ctx.get("chain_of_thought") or []

    thoughts, result = await generate(ctx, req.generator_prompt, req.task)
    memory.append(result)
    ctx.set("memory", memory)
    chain_of_thought.append({"thoughts": thoughts, "result": result})
    ctx.set("chain_of_thought", chain_of_thought)

    # Respond to the user
    return result, chain_of_thought

This lets you durably implement human control over the decision made by the LLM.
This could be extended to include tool calls when the user approves their execution. This would give you a basic version of an agent loop with explicit user-control at each step.

The benefits of using Restate here are:

What was listed for the previous example
Session/context management: you can store the session state (input items) in Restate and have it consistent and durable across calls and failures. The state can also be queried from the outside and viewed in the Restate UI.
Explicit control over what the agent should do next. Via a stateful session that is long-lived and can be resumed at any time.

Using Restate together with Agent SDKs for scalable, resilient agent loops

In this example, we show how to use Restate together with the Agent SDKs to build a scalable, resilient agent sessions.

Whereas in the first examples we called the LLM directly, in this example we use the OpenAI Agents SDK to run an agent loop for us. And we use Restate's stateful entities called Virtual Objects to manage the session state and retries for us.

The code (see this file for latest version):

agent_session = restate.VirtualObject("AgentSession")

# K/V stored in Restate
INPUT_ITEMS = "input_items"

@agent_session.handler()
async def run(ctx: restate.ObjectContext, req: InputItem) -> InputItem:
    input_items = await ctx.get(INPUT_ITEMS, type_hint=InputHistory) or InputHistory()
    input_items.entries.append(req)
    ctx.set(INPUT_ITEMS, input_items)

    async def run_agent_session() -> str:
        result = await agents.Runner.run(triage_agent, input_items.model_dump_json())
        return result.final_output
    output = await ctx.run("run agent session", run_agent_session)

    # Create a new message with the system role
    new_input_item = InputItem(role="system", content=output)
    input_items.entries.append(new_input_item)
    ctx.set(INPUT_ITEMS, input_items)
    return new_input_item

What you get from using Restate in combination with an Agent SDK:

Scalability: you can run multiple agent sessions in parallel, and Restate will manage the state for you. It will make sure a single session processes requests in order, and that the state is consistent.
Session/context management: you can store the session state (input items) in Restate and have it consistent and durable across calls and failures. The state can also be queried from the outside and viewed in the Restate UI.
Resiliency: retries and recovery of ongoing progress and state (see caveats below).
Idempotency for invocations: deduplicating requests to the same agent.

What you are missing out on if you don't integrate Restate into the agent loop (see below):

You don't log/persist any of the operations that happen within the agent loop.
The execution of an end-to-end agent loop is seen as a single non-deterministic ctx.run operation.
When something fails halfway through the loop, the loop is retried from the beginning, potentially leading to a different execution path and result.
No fine-grained retries of tool executions.
You cannot use the Restate context directly in any of the tools. Meaning, no access to Durable Execution, durable RPC, timers, etc.
From within a tool, you could execute an HTTP call to a Restate handler, but this could get scheduled multiple times on retries. It's also seen as a new invocation, so the observability experience is less nice.
Since the agent loop and tool calls do not take part in Durable Execution, it's also harder (although not impossible) to implement and execute rollback patterns when needed.

Using Restate for agentic workflows

If you look at Agent SDKs today, the most important logic is in the agent loop. In essence, this is implemented as follows:

A request (often in human language) gets forwarded to an LLM, together with a prompt describing the agent description/persona the LLM should impersonate, and which tasks it can order.
The LLM then answers with a response. This can either be a human language response that needs to be sent back to the caller, or a structured format like JSON that describes a routing decision to:
1. Another agent: In this case, we change the agent the LLM should assume (basically changing the prompt), and the respective tools it can access, and loop back to point one.
2. A tool: Execute the tool and go back to point one. The request that gets send to the LLM now includes the tool output.

Instead of using an Agent SDK, we can implement this loop with Restate constructs to get avoid the caveats we mentioned in the previous example.

This gives your agentic workflow the same resiliency and capabilities as what you get with a "traditional" workflow.

If we implement the agent loop with Restate, Restate journals each of the decisions the agents make and manages the tool executions.

The agent composes the workflow on the fly, and Restate persists the execution as it takes place.

The agent session is a Restate Virtual Object that has a handler that runs the agent loop.

The benefits of using Restate for agentic workflows are:

Flexibility and control over the loop implementation: The agent loop we implemented here can be a starting point for resilient agent sessions, that can be adapted in any way desired for the use case: e.g. more control/feedback loops over decisions, etc.
Fine-grained retries and recovery: Restate retries only the failed step and preserves the rest of the progress. This makes the system robust and fast at recovery.
Observability across agents, tools and other code execution: Restate tracks the progress requests make in a journal, that includes the entire path through the agent loop, the tool execution, and any other code or workflows that get executed as part of the request. As opposed to traces with agent SDKs, the journal is not sampled and includes the entire execution across services and tools.
Long-running workflows: Workflows can run for days, weeks, or even months. Restate persists them durably and can resume them at any time. For example, durable human-in-the loop control flow or waiting on some event to happen before continuing.
Scheduled operation: Agents can schedule tasks for later on, for example "send me an invoice after my flight". Restate tracks timers durably.
Human-in-the-loop: Restate gives you Durable Promises which make it easy to implement things like human feedback loops or external processes notifying the agent when something has happened/completed.
Task management: You can later on cancel tasks with guaranteed rollback, or latch on to retrieve their results.
Exactly-once tool calls: Restate deduplicates tool calls during retries.
Robust parallel tool calls: Restate manages the execution of the parallel tool calls and retries them if they fail.
Serverless/FaaS deployment of agent sessions and tools (e.g. AWS Lambda): Restate lets your functions suspend when they are waiting for a long time. You can mix and match: slow processes like the LLM call can run on a long-running server, while the rest of the workflow can run on serverless for demand-based scaling.

For example, we can implement an agent that can answer questions about our bank account and transactions. In this case, the agent session connected to a chat session and to the account Virtual Object:

In this example, you specify the agents, their descriptions and their tools in a configuration similar to the OpenAI Agents SDK. But in this case, all the tools are Restate handlers:

account_manager_agent = Agent(
    name="Account Manager Agent",
    handoff_description="A helpful agent that can helps you with answering questions about your bank account: the balance and transaction history.",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}
    You are an agent that helps with answering questions about the bank account: related to the balance and transaction history.
    You are not able to help with anything else.
    If you are speaking to a customer, you probably were transferred to from the intake agent.
    You have the tools to retrieve the transaction history and the balance.
    Use the following routine to support the customer.
    # Routine #1
    1. Make sure you know the customer ID. This will be the key for all interaction with the tools. 
    2. Use the get_balance and get_transaction_history tools to retrieve the balance or the transaction history of the customer.
    Depending on the customer question, use the right tool.
    3. Analyze the response and return the right information to the customer in a kind, polite, formal message.
    4. If the customer asks a question that is not related to these routines, transfer back to the intake agent.
    """,
    tools=[
        restate_tool(get_balance),
        restate_tool(get_transaction_history),
    ],
)

Here is what a slightly simplified version of the agent loop with Restate looks like. To see the full code, check out the agent_session.py file.

while True:
    # Call the LLM - OpenAPI Responses API
    response: Response = await ctx.run(
        "Call LLM",
        lambda: client.responses.create(
            model="gpt-4o",
            instructions=agent.instructions,
            input=session_state.get_input_items(),
            tools=tools_and_handoffs,
        ),
    )

    # Register the output in the session state
    session_state.add_system_messages(ctx, response.output)

    # Parse LLM response
    output_messages, run_handoffs, tool_calls = await parse_llm_response(
        ctx, agents_dict, response.output, session_state, tools
    )

    # Execute (parallel) tool calls
    parallel_tools = []
    for tool_call in tool_calls:
        if tool_call.delay_in_millis is None:
            handle = ctx.generic_call(
                service=tool_call.tool.service_name,
                handler=tool_call.tool.name,
                arg=tool_call.input_bytes,
                key=tool_call.key,
            )
            parallel_tools.append(handle)
        else:
            # Used for scheduling tasks in the future or long-running tasks like workflows
            ctx.generic_send(
                service=tool_call.tool.service_name,
                handler=tool_call.tool.name,
                arg=tool_call.input_bytes,
                key=tool_call.key,
                send_delay=timedelta(milliseconds=tool_call.delay_in_millis),
            )

    if len(parallel_tools) > 0:
        results_done = await restate.gather(*parallel_tools)
        results = [(await result).decode() for result in results_done]
        session_state.add_system_messages(ctx, results)

    # Handle handoffs
    if len(run_handoffs) > 0:
        # Only one agent can be in charge of the conversation at a time.
        # So if there are multiple handoffs in the response, only run the first one.
        handoff_command = run_handoffs[0]
        ctx.set("agent_name", handoff_command.name)

        # Start a new agent loop with the new agent
        continue

    # Handle output messages
    # If there are no output messages, then we just continue the loop
    if len(output_messages) > 0:
        return AgentResponse(
            agent=agent.name,
            messages=session_state.get_new_items(),
            final_output=output_messages[-1].content[-1].text,
        )

This code can be extended and changed based on the use case.

Run the example

Mixing static, code-defined workflows with agentic workflows

Note: We didn't need to change anything in the agent loop to make this work.

The agent session we implemented in the previous section is just a Restate Virtual Object. It can be called from anywhere, also from a more traditional code-defined workflow. For example, imagine a loan approval workflow where a step in the workflow is to wait on an agent to analyze the loan application and interact with the customer to request additional information if necessary.

Benefits of Restate here:

Benefits of previous section.
A single workflow orchestrator that handles both agentic and traditional workflows and gives the same resiliency guarantees and observability across both.

An example workflow in detail (code):

Or in code:

loan_review_workflow = restate.Workflow("LoanApprovalWorkflow")


@loan_review_workflow.main()
async def run(
    ctx: restate.WorkflowContext, loan_review_request: LoanReviewRequest
) -> LoanDecision:
    """
    Run the loan approval workflow.

    Args:
        req (LoanRequest): The loan request object.
    """
    req = loan_review_request.loan_request
    # 1. Assess the loan request
    if req.loan_amount < 1000000:
        # Loans under 1M can be approved by the AI agent
        history = loan_review_request.transaction_history.model_dump_json()
        loan_intake_data = (
            f"Loan approval workflow ID: {ctx.key()} --> Use this key to communicate back to the workflow."
            f"Review the loan request: {req.model_dump_json()};"
            f"transaction history: {history};"
        )
        from loan_review_agent import loan_review_agent
        ctx.object_send(
            agent_session_run,
            key=ctx.key(),
            arg=AgentInput(
                starting_agent=loan_review_agent,
                agents=[loan_review_agent],
                message=loan_intake_data,
            )
        )
    else:
        # Loans over 1M require human assessment
        await ctx.run("Request human assessment", request_human_assessment)

    # 2. Wait for the loan decision and add the info to the customer account
    decision = await ctx.promise("loan_decision", serde=LoanDecisionSerde).value()
    ctx.object_send(process_loan_decision, key=req.customer_id, arg=decision)

    if not decision.approved:
        return decision

    # 3. Deposit the loan amount to the customer account
    loan_transfer = Transaction(
        reason=f"Loan transfer {req.customer_id}",
        amount=req.loan_amount,
        timestamp=await time_now_string(ctx),
    )
    await ctx.object_call(account_withdraw, key="TrustworthyBank", arg=loan_transfer)
    await ctx.object_call(account_deposit, key=req.customer_id, arg=loan_transfer)

    # 4. Schedule recurring loan payments to start in one month
    ctx.object_send(
        account_loan_payment,
        key=req.customer_id,
        arg=ctx.key(),
        send_delay=timedelta(days=30),
    )
    return decision


@loan_review_workflow.handler()
async def on_loan_decision(ctx: restate.WorkflowSharedContext, decision: LoanDecision):
    """
    This tool lets you approve or reject a loan.

    Args:
        decision (LoanDecision): The result of the loan approval decision.
    """
    await ctx.promise("loan_decision", serde=LoanDecisionSerde).resolve(decision)

Run the example

Long-lived multi-agent setups

Note: We didn't need to change anything in the agent loop to make this work.

If you combine Restate Virtual Objects with Restate's scheduling capabilities, then you can build long-lived, autonomous, proactive agents which periodically come to live, for example, because of a trigger or a timer. Then, they do some tasks and possibly schedule follow-up tasks for later.

Different agent sessions could communicate with each other and distribute work among them. Each agent being its own Virtual Object.

For example, you could have a crew of specialized agents continuously monitoring social media or marketing campaigns and taking action based on what happens.

Restate would make this setup resilient, consistent, scalable, observable, and would let you run it on serverless infrastructure.

Running the examples

For all examples, you need to export your OPENAI API key as an environment variable:

export OPENAI_API_KEY=your_openai_api_key

Create a venv and install the requirements file:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt