State Rendering

using System.Text.Json.Serialization;

public class SearchInfo
{
    [JsonPropertyName("query")]
    public string Query { get; set; } = string.Empty;

    [JsonPropertyName("done")]
    public bool Done { get; set; }
}

public class AgentStateSnapshot
{
    [JsonPropertyName("searches")]
    public List<SearchInfo> Searches { get; set; } = new();
}

from typing import Annotated
from pydantic import BaseModel, Field

class SearchItem(BaseModel):
    query: str
    done: bool

# JSON schema used by AG-UI to validate and forward state to the frontend
STATE_SCHEMA: dict[str, object] = {
    "searches": {
        "type": "array",
        "items": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
                "done": {"type": "boolean"},
            },
            "required": ["query", "done"],
            "additionalProperties": False,
        },
        "description": "List of searches and whether each is done.",
    }
}

using System.Runtime.CompilerServices;
using System.Text.Json;
using AGUI.Abstractions;
using AGUI.Server;
using Microsoft.Agents.AI;
using Microsoft.Agents.AI.Hosting.AGUI.AspNetCore;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http.Json;
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Options;
using OpenAI;
using OpenAI.Chat;
using AIChatMessage = Microsoft.Extensions.AI.ChatMessage;
using AIChatResponseFormat = Microsoft.Extensions.AI.ChatResponseFormat;

var builder = WebApplication.CreateBuilder(args);
builder.Services.AddAGUIServer();
var app = builder.Build();

string openAiApiKey = builder.Configuration["OPENAI_API_KEY"]
    ?? throw new InvalidOperationException("Set OPENAI_API_KEY");

// Get JSON serializer options
var jsonOptions = app.Services.GetRequiredService<IOptions<JsonOptions>>();

// Create the base agent
AIAgent baseAgent = new OpenAIClient(openAiApiKey)
    .GetChatClient("gpt-5.4-mini")
    .AsAIAgent(
        name: "ResearchAssistant",
        instructions: "You are a research assistant that tracks your progress.");

// Wrap with state-streaming agent
AIAgent agent = new StateStreamingAgent(baseAgent, jsonOptions.Value.SerializerOptions);

// Map the AG-UI endpoint
app.MapAGUIServer("/", agent);
await app.RunAsync();

// Agent wrapper that streams state updates
internal sealed class StateStreamingAgent : DelegatingAIAgent
{
    private readonly JsonSerializerOptions _jsonSerializerOptions;

    public StateStreamingAgent(AIAgent innerAgent, JsonSerializerOptions jsonSerializerOptions)
        : base(innerAgent)
    {
        this._jsonSerializerOptions = jsonSerializerOptions;
    }

    protected override Task<AgentResponse> RunCoreAsync(
        IEnumerable<AIChatMessage> messages,
        AgentSession? session = null,
        AgentRunOptions? options = null,
        CancellationToken cancellationToken = default) =>
        RunCoreStreamingAsync(messages, session, options, cancellationToken)
            .ToAgentResponseAsync(cancellationToken);

    protected override async IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
        IEnumerable<AIChatMessage> messages,
        AgentSession? session = null,
        AgentRunOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        // Recover the state sent in the AG-UI request
        JsonElement? currentState = null;
        if (options is ChatClientAgentRunOptions { ChatOptions: { } chatOptions } &&
            chatOptions.TryGetRunAgentInput(out RunAgentInput? input) &&
            input?.State is { ValueKind: JsonValueKind.Object } state)
        {
            currentState = state;
        }

        // Clone the run options, then request structured state output
        var chatRunOptions = options as ChatClientAgentRunOptions;
        var stateChatOptions = chatRunOptions?.ChatOptions?.Clone() ?? new ChatOptions();
        stateChatOptions.ResponseFormat = AIChatResponseFormat.ForJsonSchema<AgentStateSnapshot>(
            schemaName: "AgentStateSnapshot",
            schemaDescription: "Research progress state");

        ChatClientAgentRunOptions stateOptions = new(stateChatOptions)
        {
            ChatClientFactory = chatRunOptions?.ChatClientFactory,
        };

        // Add system message with current state
        var stateMessage = new AIChatMessage(ChatRole.System,
            $"Current state: {currentState?.GetRawText() ?? "{}"}");
        var messagesWithState = messages.Append(stateMessage);

        // Collect all updates
        var allUpdates = new List<AgentResponseUpdate>();
        await foreach (var update in InnerAgent.RunStreamingAsync(messagesWithState, session, stateOptions, cancellationToken))
        {
            allUpdates.Add(update);
            // Stream non-text updates immediately
            if (update.Contents.Any(c => c is not TextContent))
            {
                yield return update;
            }
        }

        // Deserialize state snapshot from response
        var response = allUpdates.ToAgentResponse();
        JsonElement? stateSnapshot = null;
        try
        {
            stateSnapshot = JsonSerializer.Deserialize<JsonElement>(
                response.Text,
                this._jsonSerializerOptions);
        }
        catch (JsonException)
        {
            // The model did not return the requested state snapshot.
        }

        if (stateSnapshot is { } parsedStateSnapshot)
        {
            byte[] stateBytes = JsonSerializer.SerializeToUtf8Bytes(
                parsedStateSnapshot,
                this._jsonSerializerOptions.GetTypeInfo(typeof(JsonElement)));

            // Emit state snapshot as DataContent
            yield return new AgentResponseUpdate
            {
                Contents = [new DataContent(stateBytes, "application/json")]
            };
        }

        // Stream text summary
        var summaryMessage = new AIChatMessage(ChatRole.System, "Provide a brief summary of your progress.");
        await foreach (var update in InnerAgent.RunStreamingAsync(
            messages.Concat(response.Messages).Append(summaryMessage), session, options, cancellationToken))
        {
            yield return update;
        }
    }
}

from __future__ import annotations
import os
import uvicorn
from agent_framework import Agent, tool, SupportsChatGetResponse
from agent_framework.openai import OpenAIChatClient
from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
from agent_framework.ag_ui import AgentFrameworkAgent
from azure.identity import DefaultAzureCredential
from dotenv import load_dotenv
from fastapi import FastAPI
from typing import Annotated
from pydantic import BaseModel, Field

load_dotenv()

class SearchItem(BaseModel):
    query: str
    done: bool
    
STATE_SCHEMA: dict[str, object] = {
    "searches": {
        "type": "array",
        "items": {
            "type": "object",
            "properties": {
                "query": {"type": "string"},
                "done": {"type": "boolean"},
            },
            "required": ["query", "done"],
            "additionalProperties": False,
        },
        "description": "List of searches and whether each is done.",
    }
}
PREDICT_STATE_CONFIG: dict[str, dict[str, str]] = {
    "searches": {
        "tool": "update_searches",
        "tool_argument": "searches",
    }
}

@tool
def update_searches(
    searches: Annotated[list[SearchItem], Field(description=("The complete source of truth for the user's searches. Maintain ordering and include the full list on each call."))],
) -> str:
    return f"Searches updated. Tracking {len(searches)} item(s)."


def _build_chat_client():
    if os.getenv("AZURE_OPENAI_ENDPOINT"):
        azure_api_key = os.getenv("AZURE_OPENAI_API_KEY")
        return OpenAIChatClient(
            model=os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "gpt-4o-mini"),
            api_key=azure_api_key,
            credential=None if azure_api_key else DefaultAzureCredential(),
            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        )
    if os.getenv("OPENAI_API_KEY"):
        return OpenAIChatClient(
            model=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"),
            api_key=os.getenv("OPENAI_API_KEY"),
        )
    raise RuntimeError(
        "Set AZURE_OPENAI_ENDPOINT (uses az login unless AZURE_OPENAI_API_KEY is set) or OPENAI_API_KEY."
    )



def create_agent(chat_client: SupportsChatGetResponse) -> AgentFrameworkAgent:
    base_agent = Agent(
        name="search_agent",
        instructions=(
            "You help users create and run searches.\\n\\n"
            "State sync rules:\\n"
            "- Maintain a list of searches: each item has { query, done }.\\n"
            "- When adding a new search, call `update_searches` with the FULL list, including the new item with done=true.\\n"
            "- All searches in the list should have done=true unless explicitly in progress.\\n"
            "- Never send partial updates. Always include the full list on each call.\\n"
        ),
        client=chat_client,
        tools=[update_searches],
    )

    return AgentFrameworkAgent(
        agent=base_agent,
        name="CopilotKitMicrosoftAgentFrameworkAgent",
        description="Maintains a list of searches and streams state to the UI.",
        state_schema=STATE_SCHEMA,
        predict_state_config=PREDICT_STATE_CONFIG,
        require_confirmation=False,
    )
chat_client = _build_chat_client()

agent = create_agent(chat_client)

app = FastAPI(title="Microsoft Agent Framework - Quickstart")
add_agent_framework_fastapi_endpoint(app=app, agent=agent, path="/")

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)