Gestion de l’état avec AG-UI

L’état d’AG-UI est un JSON visible côté client associé à une exécution. Dans .NET, l’intégration fournit deux mécanismes explicites :

• État de lecture fourni par le client à partir de l’origine RunAgentInput.
• Associez les appels d’outils ou les résultats sélectionnés aux événements d’état d’AG-UI avec AGUIStreamOptions.

Le mappage d’état est facultatif. Les résultats de l’outil arbitraire ne deviennent pas automatiquement l’état partagé.

Lire l’état du client

MapAGUIServer stocke RunAgentInput d’origine sur ChatOptions. Si le modèle a besoin de l’état actuel du client, encapsulez l’agent de base avec un DelegatingAIAgent léger qui récupère l’état du client avec TryGetRunAgentInput et l’ajoute au contexte du modèle :

using System.Text.Json;
using AGUI.Abstractions;
using AGUI.Server;
using Microsoft.Agents.AI;
using Microsoft.Extensions.AI;

internal sealed class RecipeStateAgent(AIAgent innerAgent)
    : DelegatingAIAgent(innerAgent)
{
    protected override Task<AgentResponse> RunCoreAsync(
        IEnumerable<ChatMessage> messages,
        AgentSession? session = null,
        AgentRunOptions? options = null,
        CancellationToken cancellationToken = default) =>
        RunCoreStreamingAsync(messages, session, options, cancellationToken)
            .ToAgentResponseAsync(cancellationToken);

    protected override IAsyncEnumerable<AgentResponseUpdate> RunCoreStreamingAsync(
        IEnumerable<ChatMessage> messages,
        AgentSession? session = null,
        AgentRunOptions? options = null,
        CancellationToken cancellationToken = default)
    {
        if (options is ChatClientAgentRunOptions { ChatOptions: { } chatOptions } &&
            chatOptions.TryGetRunAgentInput(out RunAgentInput? input) &&
            input.State is { ValueKind: JsonValueKind.Object } state)
        {
            ChatMessage stateMessage = new(
                ChatRole.System,
                $"The user's current recipe state is:\n{state.GetRawText()}");
            messages = [stateMessage, .. messages];
        }

        return InnerAgent.RunStreamingAsync(
            messages,
            session,
            options,
            cancellationToken);
    }
}

AIAgent agent = new RecipeStateAgent(baseAgent);

Le wrapper gère uniquement le chemin d’accès d’entrée. L’émission d’événements d’état reste déclarative via AGUIStreamOptions, comme indiqué dans les sections suivantes. TryGetRunAgentInput lit l’entrée que la couche d’hébergement a stockée dans ChatOptions.AdditionalProperties ; le code d’application n’accède pas directement à ce dictionnaire.

L’état du client constitue des données d’entrée de la requête non fiables. Validez sa forme et ses valeurs avant de l’utiliser dans les prompts, le routage ou les opérations à privilèges.

Émettre un instantané d’état

Associer le résultat d’un outil à STATE_SNAPSHOT lorsque l’outil renvoie l’état complet :

using AGUI.Server;

AGUIStreamOptions streamOptions = new AGUIStreamOptions()
    .MapResultAsStateSnapshot("generate_recipe");

app.MapAGUIServer("/", agent).WithMetadata(streamOptions);

MapResultAsStateSnapshot nécessite que la FunctionResultContent.Result valeur soit un JsonElement. Sérialisez un POCO, un dictionnaire ou une collection en JsonElement dans l’outil avant de le renvoyer. Le résultat de generate_recipe devient alors l’instantané et remplace l’état partagé actuel du client.

Pour les autres types de résultats, utilisez MapResult avec une fonction de mappage personnalisée qui construit le StateSnapshotEvent.

Émettre des deltas d’état

Mappez un résultat d’un outil à STATE_DELTA s’il renvoie un RFC 6902 JSON Patch :

AGUIStreamOptions streamOptions = new AGUIStreamOptions()
    .MapResultAsStateSnapshot("create_plan")
    .MapResultAsStateDelta("update_plan_step");

app.MapAGUIServer("/", agent).WithMetadata(streamOptions);

Utilisez une capture instantanée pour initialiser ou remplacer l’état, et des deltas pour les modifications incrémentielles.

MapResultAsStateDelta nécessite également un JsonElement résultat. L’élément doit contenir un tableau de RFC 6902 JSON Patch. Utiliser MapResult avec un mappeur personnalisé si l’outil retourne une autre représentation.

Associer les appels d’outils à l’état

AGUIStreamOptions.MapCall associe un FunctionCallContent sélectionné à des événements AG-UI supplémentaires émis après les événements normaux d’appel d’outil. Utilisez-le lorsque l’état dérive des arguments de l’outil plutôt que du résultat de l’outil :

AGUIStreamOptions streamOptions = new AGUIStreamOptions()
    .MapCall("write_document", call =>
    {
        if (call.Arguments?.TryGetValue("document", out object? document) is not true)
        {
            return [];
        }

        JsonElement snapshot = JsonSerializer.SerializeToElement(new { document });
        return [new StateSnapshotEvent { Snapshot = snapshot }];
    });

app.MapAGUIServer("/", agent).WithMetadata(streamOptions);

L’application possède le mappage et la forme d’état. MapCall n’déduit pas l’état des arguments d’outil arbitraires ou ne supprime pas l’exécution normale de l’outil. Les mises à jour incrémentielles nécessitent que le client de modèle sous-jacent expose les arguments d’appel d’outils en continu et que l’application configure l’extraction d’argument correspondante.

État de réception dans un client .NET

Le client AG-UI .NET expose les événements de protocole d’état via ChatResponseUpdate.RawRepresentation:

await foreach (AgentResponseUpdate update in agent.RunStreamingAsync(messages, session))
{
    if (update.AsChatResponseUpdate().RawRepresentation is StateSnapshotEvent snapshot)
    {
        JsonElement state = snapshot.Snapshot;
    }
    else if (update.AsChatResponseUpdate().RawRepresentation is StateDeltaEvent delta)
    {
        JsonElement changes = delta.Delta;
    }
}

Le client est chargé de conserver et d’appliquer l’état partagé, puis d’envoyer l’état actuel sur les demandes ultérieures lorsque l’application l’exige.

Étapes suivantes

Définir des modèles d’état

Tout d’abord, définissez des modèles Pydantic pour votre structure d’état. Cela garantit la sécurité et la validation des types :

from enum import Enum
from pydantic import BaseModel, Field


class SkillLevel(str, Enum):
    """The skill level required for the recipe."""
    BEGINNER = "Beginner"
    INTERMEDIATE = "Intermediate"
    ADVANCED = "Advanced"


class CookingTime(str, Enum):
    """The cooking time of the recipe."""
    FIVE_MIN = "5 min"
    FIFTEEN_MIN = "15 min"
    THIRTY_MIN = "30 min"
    FORTY_FIVE_MIN = "45 min"
    SIXTY_PLUS_MIN = "60+ min"


class Ingredient(BaseModel):
    """An ingredient with its details."""
    icon: str = Field(..., description="Emoji icon representing the ingredient (e.g., 🥕)")
    name: str = Field(..., description="Name of the ingredient")
    amount: str = Field(..., description="Amount or quantity of the ingredient")


class Recipe(BaseModel):
    """A complete recipe."""
    title: str = Field(..., description="The title of the recipe")
    skill_level: SkillLevel = Field(..., description="The skill level required")
    special_preferences: list[str] = Field(
        default_factory=list, description="Dietary preferences (e.g., Vegetarian, Gluten-free)"
    )
    cooking_time: CookingTime = Field(..., description="The estimated cooking time")
    ingredients: list[Ingredient] = Field(..., description="Complete list of ingredients")
    instructions: list[str] = Field(..., description="Step-by-step cooking instructions")

Schéma d’état

Définissez un schéma d’état pour spécifier la structure et les types de votre état :

state_schema = {
    "recipe": {"type": "object", "description": "The current recipe"},
}

Mises à jour d’état prédictif

Les arguments de mise à jour d'état prédictive sont transmis à l'état pendant que le LLM les génère, ce qui permet des mises à jour optimistes de l'interface utilisateur.

predict_state_config = {
    "recipe": {"tool": "update_recipe", "tool_argument": "recipe"},
}

Cette configuration mappe le recipe champ d’état à l’argument recipe de l’outil update_recipe . Lorsque l’agent appelle l’outil, les arguments sont transmis à l’état en temps réel, car le LLM les génère.

Définir l’outil de mise à jour d’état

Créez une fonction d’outil qui accepte votre modèle Pydantic :

from agent_framework import tool


@tool
def update_recipe(recipe: Recipe) -> str:
    """Update the recipe with new or modified content.

    You MUST write the complete recipe with ALL fields, even when changing only a few items.
    When modifying an existing recipe, include ALL existing ingredients and instructions plus your changes.
    NEVER delete existing data - only add or modify.

    Args:
        recipe: The complete recipe object with all details

    Returns:
        Confirmation that the recipe was updated
    """
    return "Recipe updated."

Créer l’agent avec State Management

Voici une implémentation complète du serveur avec la gestion de l’état :

"""AG-UI server with state management."""

from agent_framework import Agent
from agent_framework.openai import OpenAIChatCompletionClient
from agent_framework_ag_ui import (
    AgentFrameworkAgent,
    add_agent_framework_fastapi_endpoint,
)
from azure.identity import AzureCliCredential
from fastapi import FastAPI

# Create the chat agent with tools
agent = Agent(
    name="recipe_agent",
    instructions="""You are a helpful recipe assistant that creates and modifies recipes.

    CRITICAL RULES:
    1. You will receive the current recipe state in the system context
    2. To update the recipe, you MUST use the update_recipe tool
    3. When modifying a recipe, ALWAYS include ALL existing data plus your changes in the tool call
    4. NEVER delete existing ingredients or instructions - only add or modify
    5. After calling the tool, provide a brief conversational message (1-2 sentences)

    When creating a NEW recipe:
    - Provide all required fields: title, skill_level, cooking_time, ingredients, instructions
    - Use actual emojis for ingredient icons (🥕 🧄 🧅 🍅 🌿 🍗 🥩 🧀)
    - Leave special_preferences empty unless specified
    - Message: "Here's your recipe!" or similar

    When MODIFYING or IMPROVING an existing recipe:
    - Include ALL existing ingredients + any new ones
    - Include ALL existing instructions + any new/modified ones
    - Update other fields as needed
    - Message: Explain what you improved (e.g., "I upgraded the ingredients to premium quality")
    - When asked to "improve", enhance with:
      * Better ingredients (upgrade quality, add complementary flavors)
      * More detailed instructions
      * Professional techniques
      * Adjust skill_level if complexity changes
      * Add relevant special_preferences

    Example improvements:
    - Upgrade "chicken" → "organic free-range chicken breast"
    - Add herbs: basil, oregano, thyme
    - Add aromatics: garlic, shallots
    - Add finishing touches: lemon zest, fresh parsley
    - Make instructions more detailed and professional
    """,
    client=OpenAIChatCompletionClient(
        model=deployment_name,
        azure_endpoint=endpoint,
        api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
        credential=AzureCliCredential(),
    ),
    tools=[update_recipe],
)

# Wrap agent with state management
recipe_agent = AgentFrameworkAgent(
    agent=agent,
    name="RecipeAgent",
    description="Creates and modifies recipes with streaming state updates",
    state_schema={
        "recipe": {"type": "object", "description": "The current recipe"},
    },
    predict_state_config={
        "recipe": {"tool": "update_recipe", "tool_argument": "recipe"},
    },
)

# Create FastAPI app
app = FastAPI(title="AG-UI Recipe Assistant")
add_agent_framework_fastapi_endpoint(app, recipe_agent, "/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8888)

Concepts clés

• Modèles Pydantic : Définir un état structuré avec la sécurité et la validation des types
• Schéma d’état : format simple spécifiant les types de champs d’état
• Configuration de l’état prédictif : mappe les champs d’état aux arguments d’outil pour les mises à jour en streaming
• Injection d’état : l’état actuel est automatiquement injecté en tant que messages système pour fournir un contexte
• Mises à jour complètes : les outils doivent écrire l’état complet, pas seulement les deltas
• Stratégie de confirmation : personnaliser les messages d’approbation pour votre domaine (recette, document, planification des tâches, etc.)

Compréhension des événements d’état

Événement instantané d'état

Capture instantanée complète de l’état actuel, émise lorsque l’outil se termine :

{
    "type": "STATE_SNAPSHOT",
    "snapshot": {
        "recipe": {
            "title": "Classic Pasta Carbonara",
            "skill_level": "Intermediate",
            "special_preferences": ["Authentic Italian"],
            "cooking_time": "30 min",
            "ingredients": [
                {"icon": "🍝", "name": "Spaghetti", "amount": "400g"},
                {"icon": "🥓", "name": "Guanciale or bacon", "amount": "200g"},
                {"icon": "🥚", "name": "Egg yolks", "amount": "4"},
                {"icon": "🧀", "name": "Pecorino Romano", "amount": "100g grated"},
                {"icon": "🧂", "name": "Black pepper", "amount": "To taste"}
            ],
            "instructions": [
                "Bring a large pot of salted water to boil",
                "Cut guanciale into small strips and fry until crispy",
                "Beat egg yolks with grated Pecorino and black pepper",
                "Cook spaghetti until al dente",
                "Reserve 1 cup pasta water, then drain pasta",
                "Remove pan from heat, add hot pasta to guanciale",
                "Quickly stir in egg mixture, adding pasta water to create creamy sauce",
                "Serve immediately with extra Pecorino and black pepper"
            ]
        }
    }
}

Événement Delta d’état

Mises à jour d’état incrémentielles au format patch JSON, émises en tant qu’arguments de l’outil de flux LLM :

{
    "type": "STATE_DELTA",
    "delta": [
        {
            "op": "replace",
            "path": "/recipe",
            "value": {
                "title": "Classic Pasta Carbonara",
                "skill_level": "Intermediate",
                "cooking_time": "30 min",
                "ingredients": [
                    {"icon": "🍝", "name": "Spaghetti", "amount": "400g"}
                ],
                "instructions": ["Bring a large pot of salted water to boil"]
            }
        }
    ]
}

Mise en œuvre chez le client

Le package agent_framework_ag_ui fournit AGUIChatClient pour la connexion aux serveurs AG-UI, offrant une expérience utilisateur Python à la parité avec .NET.

"""AG-UI client with state management."""

import asyncio
import json
import os
from typing import Any

from agent_framework import Agent, Message, Role
from agent_framework_ag_ui import AGUIChatClient


async def main():
    """Example client with state tracking."""
    server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:8888/")
    print(f"Connecting to AG-UI server at: {server_url}\n")

    # Create AG-UI chat client
    chat_client = AGUIChatClient(endpoint=server_url)

    # Wrap with Agent for convenient API
    agent = Agent(
        name="ClientAgent",
        client=chat_client,
        instructions="You are a helpful assistant.",
    )

    # Get a thread for conversation continuity
    thread = agent.create_session()

    # Track state locally
    state: dict[str, Any] = {}

    try:
        while True:
            message = input("\nUser (:q to quit, :state to show state): ")
            if not message.strip():
                continue

            if message.lower() in (":q", "quit"):
                break

            if message.lower() == ":state":
                print(f"\nCurrent state: {json.dumps(state, indent=2)}")
                continue

            print()
            # Stream the agent response with state
            async for update in agent.run(message, session=thread, stream=True):
                # Handle text content
                if update.text:
                    print(update.text, end="", flush=True)

                # Handle state updates surfaced through AG-UI events.
                for content in update.contents:
                    if content.type == "data" and getattr(content, "media_type", None) == "application/json":
                        print("\n[JSON state payload received]")

            print(f"\n\nCurrent state: {json.dumps(state, indent=2)}")
            print()

    except KeyboardInterrupt:
        print("\n\nExiting...")


if __name__ == "__main__":
    # Install dependencies: pip install agent-framework-ag-ui --pre
    asyncio.run(main())

Principaux avantages

Le AGUIChatClient fournit les éléments suivants :

• Connexion simplifiée : gestion automatique de la communication HTTP/SSE
• Gestion des threads : suivi d’ID de thread intégré pour la continuité des conversations
• Intégration de l’agent : fonctionne en toute transparence avec Agent pour l’API familière
• Gestion de l’état : analyse automatique des événements d’état à partir du serveur
• Parité avec .NET : expérience cohérente entre les langages

Confirmation de l’état prédit

Définissez require_confirmation=True sur AgentFrameworkAgent lorsque les changements d’état prédits doivent attendre une confirmation du client avant d’être appliqués :

recipe_agent = AgentFrameworkAgent(
    agent=agent,
    state_schema={"recipe": {"type": "object", "description": "The current recipe"}},
    predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
    require_confirmation=True,
)

Personnalisez la copie de confirmation dans votre interface utilisateur du client AG-UI lors du rendu de l’événement de confirmation.

Exemple d’interaction

Avec le serveur et le client en cours d’exécution :

User (:q to quit, :state to show state): I want to make a classic Italian pasta carbonara

[Run Started]
[Calling Tool: update_recipe]
[State Updated]
[State Updated]
[State Updated]
[Tool Result: Recipe updated.]
Here's your recipe!
[Run Finished]

============================================================
CURRENT STATE
============================================================

recipe:
  title: Classic Pasta Carbonara
  skill_level: Intermediate
  special_preferences: ['Authentic Italian']
  cooking_time: 30 min
  ingredients:
    - 🍝 Spaghetti: 400g
    - 🥓 Guanciale or bacon: 200g
    - 🥚 Egg yolks: 4
    - 🧀 Pecorino Romano: 100g grated
    - 🧂 Black pepper: To taste
  instructions:
    1. Bring a large pot of salted water to boil
    2. Cut guanciale into small strips and fry until crispy
    3. Beat egg yolks with grated Pecorino and black pepper
    4. Cook spaghetti until al dente
    5. Reserve 1 cup pasta water, then drain pasta
    6. Remove pan from heat, add hot pasta to guanciale
    7. Quickly stir in egg mixture, adding pasta water to create creamy sauce
    8. Serve immediately with extra Pecorino and black pepper

============================================================

Mises à jour d’état prédictif en action

Lors de l’utilisation des mises à jour d’état prédictives avec predict_state_config, le client reçoit les événements STATE_DELTA lorsque le LLM génère des arguments d’outil en temps réel, avant que l’outil ne s’exécute :

// Agent starts generating tool call for update_recipe
// Client receives STATE_DELTA events as the recipe argument streams:

// First delta - partial recipe with title
{
  "type": "STATE_DELTA",
  "delta": [{"op": "replace", "path": "/recipe", "value": {"title": "Classic Pasta"}}]
}

// Second delta - title complete with more fields
{
  "type": "STATE_DELTA",
  "delta": [{"op": "replace", "path": "/recipe", "value": {
    "title": "Classic Pasta Carbonara",
    "skill_level": "Intermediate"
  }}]
}

// Third delta - ingredients starting to appear
{
  "type": "STATE_DELTA",
  "delta": [{"op": "replace", "path": "/recipe", "value": {
    "title": "Classic Pasta Carbonara",
    "skill_level": "Intermediate",
    "cooking_time": "30 min",
    "ingredients": [
      {"icon": "🍝", "name": "Spaghetti", "amount": "400g"}
    ]
  }}]
}

// ... more deltas as the LLM generates the complete recipe

Cela permet au client d’afficher les mises à jour optimistes de l’interface utilisateur en temps réel, car l’agent pense, en fournissant des commentaires immédiats aux utilisateurs.

État avec Humain dans la boucle

Vous pouvez combiner la gestion de l’état avec les flux de travail d’approbation en définissant require_confirmation=True:

recipe_agent = AgentFrameworkAgent(
    agent=agent,
    state_schema={"recipe": {"type": "object", "description": "The current recipe"}},
    predict_state_config={"recipe": {"tool": "update_recipe", "tool_argument": "recipe"}},
    require_confirmation=True,  # Require approval for state changes
)

Quand cette option est activée :

1. Mise à jour continue de l'état quand l’agent génère des arguments d'outil (mises à jour prédictives via des événements STATE_DELTA)
2. L’agent se met en pause avant d’exécuter l’outil avec un tool_call signal d’interruption dans RUN_FINISHED.outcome.interrupts
3. S’il est approuvé, l’outil s’exécute et l’état final est émis (via STATE_SNAPSHOT l’événement)
4. Si elles sont rejetées, les modifications d'état prédictives sont abandonnées.

Modèles d’état avancés

État complexe avec plusieurs champs

Vous pouvez gérer plusieurs champs d’état avec différents outils :

from pydantic import BaseModel


class TaskStep(BaseModel):
    """A single task step."""
    description: str
    status: str = "pending"
    estimated_duration: str = "5 min"


@tool
def generate_task_steps(steps: list[TaskStep]) -> str:
    """Generate task steps for a given task."""
    return f"Generated {len(steps)} steps."


@tool
def update_preferences(preferences: dict[str, Any]) -> str:
    """Update user preferences."""
    return "Preferences updated."


# Configure with multiple state fields
agent_with_multiple_state = AgentFrameworkAgent(
    agent=agent,
    state_schema={
        "steps": {"type": "array", "description": "List of task steps"},
        "preferences": {"type": "object", "description": "User preferences"},
    },
    predict_state_config={
        "steps": {"tool": "generate_task_steps", "tool_argument": "steps"},
        "preferences": {"tool": "update_preferences", "tool_argument": "preferences"},
    },
)

Utilisation d’arguments d’outil générique

Lorsqu’un outil retourne des données imbriquées complexes, utilisez "*" pour mapper tous les arguments de l’outil à l’état :

@tool
def create_document(title: str, content: str, metadata: dict[str, Any]) -> str:
    """Create a document with title, content, and metadata."""
    return "Document created."


# Map all tool arguments to document state
predict_state_config = {
    "document": {"tool": "create_document", "tool_argument": "*"}
}

Cela mappe l’ensemble de l’appel de l’outil (tous les arguments) au document champ d’état.

Meilleures pratiques

Utiliser des modèles Pydantic

Définissez des modèles structurés pour la sécurité des types :

class Recipe(BaseModel):
    """Use Pydantic models for structured, validated state."""
    title: str
    skill_level: SkillLevel
    ingredients: list[Ingredient]
    instructions: list[str]

Avantages :

• Sécurité des types de données : validation automatique des types de données
• Documentation : Les descriptions de champs servent de documentation
• Prise en charge de l’IDE : saisie semi-automatique et contrôle de type
• Sérialisation : conversion JSON automatique

Mises à jour d’état complètes

Écrivez toujours l’état complet, pas seulement les deltas :

@tool
def update_recipe(recipe: Recipe) -> str:
    """
    You MUST write the complete recipe with ALL fields.
    When modifying a recipe, include ALL existing ingredients and
    instructions plus your changes. NEVER delete existing data.
    """
    return "Recipe updated."

Cela garantit la cohérence de l’état et les mises à jour prédictives appropriées.

Mettre en correspondance les noms des paramètres

Vérifiez que les noms des paramètres de l’outil correspondent à tool_argument la configuration :

# Tool parameter name
def update_recipe(recipe: Recipe) -> str:  # Parameter name: 'recipe'
    ...

# Must match in predict_state_config
predict_state_config = {
    "recipe": {"tool": "update_recipe", "tool_argument": "recipe"}  # Same name
}

Fournir un contexte dans les instructions

Incluez des instructions claires sur la gestion de l’état :

agent = Agent(
    instructions="""
    CRITICAL RULES:
    1. You will receive the current recipe state in the system context
    2. To update the recipe, you MUST use the update_recipe tool
    3. When modifying a recipe, ALWAYS include ALL existing data plus your changes
    4. NEVER delete existing ingredients or instructions - only add or modify
    """,
    ...
)

Personnaliser l’interface utilisateur de confirmation

Personnalisez les messages d’approbation et de confirmation d’état dans votre client AG-UI lors du rendu des événements de confirmation à partir du serveur.

Prochaines étapes

Vous avez maintenant appris toutes les fonctionnalités de base AG-UI ! Vous pouvez ensuite :

• Explorer la documentation de l’Infrastructure de l’agent
• Créer une application complète combinant toutes les fonctionnalités de AG-UI
• Déployer votre service AG-UI en production

Ressources additionnelles

• Vue d’ensemble deAG-UI
• Mise en route
• Rendu de l'outil côté serveur

La gestion de l’état d’AG-UI en Go peut être mise en œuvre avec un middleware qui émet des mises à jour structurées message.DataContent en parallèle de mises à jour de texte classiques.

stateSnapshotMiddleware := agent.MiddlewareFunc(func(next agent.RunFunc, ctx context.Context, messages []*message.Message, opts ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] {
    return func(yield func(*agent.ResponseUpdate, error) bool) {
        for update, err := range next(ctx, messages, opts...) {
            if err != nil {
                yield(nil, err)
                return
            }
            if update != nil {
                // Inspect update contents and yield DataContent snapshots as needed.
            }
            if !yield(update, nil) {
                return
            }
        }
    }
})

a := foundryprovider.NewAgent(endpoint, token, foundryprovider.ModelDeployment(model), foundryprovider.AgentConfig{
    Config: agent.Config{
        Middlewares: []agent.Middleware{stateSnapshotMiddleware},
    },
})