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 Azure.AI.OpenAI;
using Azure.Identity;
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;

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

string endpoint = builder.Configuration["AZURE_OPENAI_ENDPOINT"]!;
string deployment = builder.Configuration["AZURE_OPENAI_DEPLOYMENT_NAME"]!;

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

// Create the base agent
AIAgent baseAgent = new AzureOpenAIClient(new Uri(endpoint), new DefaultAzureCredential())
    .GetChatClient(deployment)
    .CreateAIAgent(
        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.MapAGUI("/", 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;
    }

    public override async IAsyncEnumerable<AgentRunResponseUpdate> RunStreamingAsync(
        IEnumerable<ChatMessage> messages,
        AgentThread? thread = null,
        AgentRunOptions? options = null,
        [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        // Get current state from options if provided
        JsonElement currentState = default;
        if (options is ChatClientAgentRunOptions { ChatOptions.AdditionalProperties: { } properties } &&
            properties.TryGetValue("ag_ui_state", out object? stateObj) && stateObj is JsonElement state)
        {
            currentState = state;
        }

        // Create options with JSON schema for structured state output
        ChatClientAgentRunOptions stateOptions = new ChatClientAgentRunOptions
        {
            ChatOptions = new ChatOptions
            {
                ResponseFormat = ChatResponseFormat.ForJsonSchema<AgentStateSnapshot>(
                    schemaName: "AgentStateSnapshot",
                    schemaDescription: "Research progress state")
            }
        };

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

        // Collect all updates
        var allUpdates = new List<AgentRunResponseUpdate>();
        await foreach (var update in this.InnerAgent.RunStreamingAsync(messagesWithState, thread, 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.ToAgentRunResponse();
        if (response.TryDeserialize(this._jsonSerializerOptions, out JsonElement stateSnapshot))
        {
            byte[] stateBytes = JsonSerializer.SerializeToUtf8Bytes(
                stateSnapshot,
                this._jsonSerializerOptions.GetTypeInfo(typeof(JsonElement)));

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

        // Stream text summary
        var summaryMessage = new ChatMessage(ChatRole.System, "Provide a brief summary of your progress.");
        await foreach (var update in this.InnerAgent.RunStreamingAsync(
            messages.Concat(response.Messages).Append(summaryMessage), thread, 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 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"):
        return OpenAIChatClient(
            model=os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "gpt-4o-mini"),
            api_key=os.getenv("AZURE_OPENAI_API_KEY"),
            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 either AZURE_OPENAI_ENDPOINT + AZURE_OPENAI_API_KEY, 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)