assistant-ui Quickstart

import random
from typing import Annotated

from ag2 import Agent, Context, tool
from ag2.ag_ui import AGUIStream, read_run_input
from ag2.config import OpenAIResponsesConfig
from ag2.middleware.builtin.tools.approval import ApprovalRequired
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import Field

CONDITIONS = ("sunny", "partly cloudy", "overcast", "light rain", "windy")

@tool(description="Get the current weather for a city.")
def get_weather(
    city: Annotated[str, Field(description="The city name, e.g. 'Lisbon'")],
) -> dict[str, str | int]:
    # A stand-in for a real weather API, seeded by the city so answers are stable.
    rng = random.Random(city.lower())
    return {
        "city": city,
        "temperature_c": rng.randint(-5, 32),
        "conditions": rng.choice(CONDITIONS),
    }

@tool(
    description=(
        "Save a city to the user's favorites. Use only when the user asks to save, "
        "pin or remember a city. The user confirms before anything is saved."
    ),
    # The call is held as an AG-UI interrupt until the client approves or rejects it.
    middleware=[
        ApprovalRequired("Save {tool_arguments} to your favorites?", allow_always=False)
    ],
)
def save_favorite_city(
    context: Context,
    city: Annotated[str, Field(description="The city name to save, e.g. 'Lisbon'")],
) -> list[str]:
    # Variables are sent to the client as the shared state snapshot.
    favorites = [*context.variables.get("favorites", []), city]
    context.variables["favorites"] = favorites
    return favorites

agent = Agent(
    "WeatherAgent",
    prompt=(
        "You are a concise weather assistant. Use get_weather for weather questions. "
        "Use save_favorite_city only when the user asks to save a city. "
        "After a tool call, answer in one short sentence."
    ),
    config=OpenAIResponsesConfig(model="gpt-6-luna", streaming=True),
    tools=[get_weather, save_favorite_city],
)

stream = AGUIStream(agent)

app = FastAPI()
# The browser calls this server directly, so the Next.js dev server's origin is allowed.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000"],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

@app.get("/agent")
async def capabilities() -> JSONResponse:
    """Tell the client what the agent supports before a run starts."""
    return JSONResponse(
        stream.capabilities().model_dump(by_alias=True, exclude_none=True)
    )

@app.post("/agent")
async def run_agent(request: Request) -> Response:
    """Run the agent on an AG-UI `RunAgentInput` and stream its events as SSE."""
    try:
        run_input = read_run_input(await request.body())
    except ValueError:
        return JSONResponse(
            {"error": "invalid AG-UI RunAgentInput body"}, status_code=400
        )
    return StreamingResponse(
        stream.dispatch(run_input, accept=request.headers.get("accept")),
        media_type="text/event-stream",
    )

"use client";

import { useMemo, type ReactNode } from "react";
import { AssistantRuntimeProvider } from "@assistant-ui/react";
import { HttpAgent } from "@ag-ui/client";
import { useAgUiRuntime } from "@assistant-ui/react-ag-ui";

export function MyRuntimeProvider({ children }: { children: ReactNode }) {
  const agent = useMemo(
    () => new HttpAgent({ url: "http://localhost:8000/agent" }),
    [],
  );
  const runtime = useAgUiRuntime({ agent });

  return (
    <AssistantRuntimeProvider runtime={runtime}>
      {children}
    </AssistantRuntimeProvider>
  );
}

import { AuiConfig, AuiProvider, Tools, defineToolkit, useAui } from "@assistant-ui/react";

type Weather = { city: string; temperature_c: number; conditions: string };

const toolkit = defineToolkit({
  get_weather: {
    type: "backend",
    render: ({ args, result }) => {
      const weather = result as Weather | undefined;
      return weather ? (
        <p>{weather.city}: {weather.temperature_c}°C, {weather.conditions}</p>
      ) : (
        <p>Looking up {String(args["city"] ?? "…")}</p>
      );
    },
  },
});

// Inside your page component:
const aui = useAui();
const config = AuiConfig({ tools: Tools({ toolkit }) });
// <AuiProvider extends={aui} config={config}><Thread /></AuiProvider>

import {
  useAgUiInterrupts,
  useAgUiSubmitInterruptResponses,
} from "@assistant-ui/react-ag-ui";

function PendingApprovals() {
  const interrupts = useAgUiInterrupts();
  const submit = useAgUiSubmitInterruptResponses();

  const answer = (approved: boolean) =>
    submit(
      interrupts.map((i) => ({
        interruptId: i.id,
        status: "resolved" as const,
        payload: approved,
      })),
    );

  return interrupts.map((i) => (
    <div key={i.id}>
      <p>{i.message}</p>
      <button onClick={() => void answer(true)}>Approve</button>
      <button onClick={() => void answer(false)}>Deny</button>
    </div>
  ));
}

import { useAgUiState } from "@assistant-ui/react-ag-ui";

function Favorites() {
  const state = useAgUiState<{ favorites?: string[] }>();
  return <p>Favorites: {state?.favorites?.join(", ") || "none"}</p>;
}