Middleware del agente

Agent Framework se puede personalizar mediante tres tipos diferentes de middleware:

1. Middleware de ejecución del agente: permite interceptar todas las ejecuciones del agente, de modo que la entrada y la salida se puedan inspeccionar o modificar según sea necesario.
2. Middleware de llamada a funciones: permite la interceptación de todas las llamadas de función ejecutadas por el agente, de modo que la entrada y la salida se puedan inspeccionar y modificar según sea necesario.
3. IChatClient middleware: permite la interceptación de llamadas a una IChatClient implementación, donde un agente usa IChatClient para llamadas de inferencia, por ejemplo, al usar ChatClientAgent.

Todos los tipos de middleware se implementan a través de una devolución de llamada de función y cuando se registran varias instancias de middleware del mismo tipo, forman una cadena, donde se espera que cada instancia de middleware llame a la siguiente de la cadena, a través de un proporcionado nextFunc.

Los tipos de middleware de ejecución y de llamada de agente se pueden registrar en un agente mediante el generador de agentes con un objeto de agente existente.

var middlewareEnabledAgent = originalAgent
    .AsBuilder()
        .Use(runFunc: CustomAgentRunMiddleware, runStreamingFunc: CustomAgentRunStreamingMiddleware)
        .Use(CustomFunctionCallingMiddleware)
    .Build();

IChatClient el middleware se puede registrar en un IChatClient antes de que se use con , ChatClientAgentmediante el patrón del generador de clientes de chat.

var chatClient = new AIProjectClient(
    new Uri("<your-foundry-project-endpoint>"),
    new DefaultAzureCredential())
        .GetProjectOpenAIClient()
        .GetProjectResponsesClient()
        .AsIChatClient(deploymentName);

var middlewareEnabledChatClient = chatClient
    .AsBuilder()
        .Use(getResponseFunc: CustomChatClientMiddleware, getStreamingResponseFunc: null)
    .Build();

var agent = new ChatClientAgent(middlewareEnabledChatClient, instructions: "You are a helpful assistant.");

IChatClient el middleware también se puede registrar mediante un método de fábrica al construir un agente a través de uno de los métodos auxiliares en los clientes del SDK.

var agent = new AIProjectClient(
    new Uri("<your-foundry-project-endpoint>"),
    new DefaultAzureCredential())
        .AsAIAgent(
            model: deploymentName,
            instructions: "You are a helpful assistant.",
            clientFactory: (chatClient) => chatClient
                .AsBuilder()
                    .Use(getResponseFunc: CustomChatClientMiddleware, getStreamingResponseFunc: null)
                .Build());

Middleware de ejecución del agente

Este es un ejemplo de middleware de ejecución del agente, que puede inspeccionar o modificar la entrada y salida de la ejecución del agente.

async Task<AgentResponse> CustomAgentRunMiddleware(
    IEnumerable<ChatMessage> messages,
    AgentSession? session,
    AgentRunOptions? options,
    AIAgent innerAgent,
    CancellationToken cancellationToken)
{
    Console.WriteLine(messages.Count());
    var response = await innerAgent.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false);
    Console.WriteLine(response.Messages.Count);
    return response;
}

Middleware de ejecución de streaming del agente

Este es un ejemplo de middleware de streaming de ejecución de agente, que puede inspeccionar o modificar la entrada y salida de la ejecución de streaming del agente.

    async IAsyncEnumerable<AgentResponseUpdate> CustomAgentRunStreamingMiddleware(
    IEnumerable<ChatMessage> messages,
    AgentSession? session,
    AgentRunOptions? options,
    AIAgent innerAgent,
    [EnumeratorCancellation] CancellationToken cancellationToken)
{
    Console.WriteLine(messages.Count());
    List<AgentResponseUpdate> updates = [];
    await foreach (var update in innerAgent.RunStreamingAsync(messages, session, options, cancellationToken))
    {
        updates.Add(update);
        yield return update;
    }

    Console.WriteLine(updates.ToAgentResponse().Messages.Count);
}

Middleware de llamada a funciones

Este es un ejemplo de un middleware que llama a una función, el cual puede inspeccionar y/o modificar la función que se llama y el resultado de dicha llamada.

async ValueTask<object?> CustomFunctionCallingMiddleware(
    AIAgent agent,
    FunctionInvocationContext context,
    Func<FunctionInvocationContext, CancellationToken, ValueTask<object?>> next,
    CancellationToken cancellationToken)
{
    Console.WriteLine($"Function Name: {context!.Function.Name}");
    var result = await next(context, cancellationToken);
    Console.WriteLine($"Function Call Result: {result}");

    return result;
}

Es posible finalizar el bucle de llamada de función con middleware de llamada de función estableciendo el proporcionado FunctionInvocationContext.Terminate en true. Esto impedirá que el bucle de llamada de función emita una solicitud al servicio de inferencia que contiene los resultados de la llamada de función después de la invocación de función. Si había más de una función disponible para la invocación durante esta iteración, también podría impedir que se ejecuten las funciones restantes.

Middleware IChatClient

Este es un ejemplo de middleware de cliente de chat, que puede inspeccionar o modificar la entrada y salida de la solicitud al servicio de inferencia que proporciona el cliente de chat.

async Task<ChatResponse> CustomChatClientMiddleware(
    IEnumerable<ChatMessage> messages,
    ChatOptions? options,
    IChatClient innerChatClient,
    CancellationToken cancellationToken)
{
    Console.WriteLine(messages.Count());
    var response = await innerChatClient.GetResponseAsync(messages, options, cancellationToken);
    Console.WriteLine(response.Messages.Count);

    return response;
}

Agent Framework se puede personalizar mediante tres tipos diferentes de middleware:

1. Middleware del agente: intercepta la ejecución de la ejecución del agente, lo que le permite inspeccionar y modificar entradas, salidas y flujo de control.
2. Middleware de función: intercepta las llamadas de función (herramienta) realizadas durante la ejecución del agente, habilitando la validación de entrada, la transformación de resultados y el control de ejecución.
3. Middleware de chat: intercepta las solicitudes de chat subyacentes enviadas a los modelos de IA, lo que proporciona acceso a los mensajes, opciones y respuestas sin procesar.

Todos los tipos admiten implementaciones basadas en funciones y basadas en clases. Cuando se registran varios middleware del mismo tipo, forman una cadena donde cada una llama a la call_next devolución de llamada para continuar el procesamiento. call_next no toma el contexto como argumento; middleware muta el objeto de contexto compartido directamente y, a continuación, espera call_next().

Middleware del agente

El middleware del agente intercepta y modifica la ejecución de la ejecución del agente. Usa el AgentContext objeto que contiene:

• agent: el agente al que se invoca.
• messages: lista de mensajes de chat en la conversación
• session: la sesión del agente actual, si existe.
• options: opciones de ejecución del agente para esta invocación
• stream: booleano que indica si la respuesta es streaming.
• metadata: diccionario para almacenar datos adicionales entre middleware
• result: la respuesta del agente (se puede modificar)
• kwargs: argumentos de palabra clave del entorno de ejecución heredados pasados al método de ejecución del agente
• client_kwargs: valores de tiempo de ejecución específicos del cliente para clientes de chat de nivel inferior
• function_invocation_kwargs: valores en tiempo de ejecución que se reenviarán a las herramientas.

La call_next devolución de llamada continúa la cadena de middleware o ejecuta el agente si es el último middleware.

Basado en funciones

async def inject_tool_runtime_defaults(
    context: AgentContext,
    call_next: Callable[[], Awaitable[None]],
) -> None:
    """Agent middleware that sets tool-only runtime defaults."""
    print("[Agent] Starting execution")
    context.function_invocation_kwargs.setdefault("tenant", "contoso")
    context.function_invocation_kwargs.setdefault("request_source", "agent-middleware")

    await call_next()

    print("[Agent] Execution completed")

Basado en clases

El middleware del agente basado en clases usa un process método que tiene la misma firma y comportamiento que el middleware basado en funciones.

from agent_framework import AgentMiddleware, AgentContext

class LoggingAgentMiddleware(AgentMiddleware):
    """Agent middleware that logs execution."""

    async def process(
        self,
        context: AgentContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        print("[Agent Class] Starting execution")
        await call_next()
        print("[Agent Class] Execution completed")

Middleware de función

El middleware de función intercepta las llamadas de función dentro de los agentes. Usa el FunctionInvocationContext objeto que contiene:

• function: la función que se invoca.
• arguments: argumentos validados para la función
• session: la sesión del agente actual, si existe.
• metadata: diccionario para almacenar datos adicionales entre middleware
• result: el valor devuelto de la función (se puede modificar)
• kwargs: argumentos de palabra clave en tiempo de ejecución que se reenviarán a la invocación de la herramienta.

La call_next devolución de llamada continúa con el siguiente middleware o ejecuta la función real.

Basado en funciones

async def inject_function_kwargs(
    context: FunctionInvocationContext,
    call_next: Callable[[], Awaitable[None]],
) -> None:
    """Function middleware that enriches tool runtime values."""
    context.kwargs.setdefault("tenant", "contoso")
    context.kwargs.setdefault("request_source", "function-middleware")

    await call_next()

Basado en clases

from agent_framework import FunctionMiddleware, FunctionInvocationContext

class LoggingFunctionMiddleware(FunctionMiddleware):
    """Function middleware that logs function execution."""

    async def process(
        self,
        context: FunctionInvocationContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        print(f"[Function Class] Calling {context.function.name}")
        await call_next()
        print(f"[Function Class] {context.function.name} completed")

Chat Middleware

El middleware de chat intercepta las solicitudes de chat enviadas a los modelos de IA. Usa el ChatContext objeto que contiene:

• chat_client: el cliente de chat que se invoca.
• messages: lista de mensajes que se envían al servicio de IA
• options: las opciones de la solicitud de chat
• stream: booleano que indica si se trata de una invocación de streaming.
• metadata: diccionario para almacenar datos adicionales entre middleware
• result: la respuesta del chat de la inteligencia artificial (se puede modificar)
• kwargs: argumentos de palabra clave adicionales pasados al cliente de chat
• function_invocation_kwargs: valores en tiempo de ejecución de solo herramienta que reenviará la capa de chat.

La call_next devolución de llamada continúa con el siguiente middleware o envía la solicitud al servicio de IA.

Basado en funciones

async def logging_chat_middleware(
    context: ChatContext,
    call_next: Callable[[], Awaitable[None]],
) -> None:
    """Chat middleware that logs AI interactions."""
    # Pre-processing: Log before AI call
    print(f"[Chat] Sending {len(context.messages)} messages to AI")

    # Continue to next middleware or AI service
    await call_next()

    # Post-processing: Log after AI response
    print("[Chat] AI response received")

Basado en clases

from agent_framework import ChatMiddleware, ChatContext

class LoggingChatMiddleware(ChatMiddleware):
    """Chat middleware that logs AI interactions."""

    async def process(
        self,
        context: ChatContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        print(f"[Chat Class] Sending {len(context.messages)} messages to AI")
        await call_next()
        print("[Chat Class] AI response received")

Decoradores de middleware

Los decoradores proporcionan una declaración de tipo de middleware explícita sin necesidad de anotaciones de tipo. Son útiles cuando no se usan anotaciones de tipo o se quieren evitar errores de coincidencia de tipos:

from agent_framework import agent_middleware, function_middleware, chat_middleware

@agent_middleware
async def simple_agent_middleware(context, call_next):
    print("Before agent execution")
    await call_next()
    print("After agent execution")

@function_middleware
async def simple_function_middleware(context, call_next):
    print(f"Calling function: {context.function.name}")
    await call_next()
    print("Function call completed")

@chat_middleware
async def simple_chat_middleware(context, call_next):
    print(f"Processing {len(context.messages)} chat messages")
    await call_next()
    print("Chat processing completed")

Registro de middleware

El middleware se puede registrar en dos niveles con distintos ámbitos y comportamientos.

Middleware de Agent-Level frente a Run-Level

from agent_framework import Agent
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential

# Agent-level middleware: Applied to ALL runs of the agent
async with Agent(
    client=FoundryChatClient(credential=credential),
    name="WeatherAgent",
    instructions="You are a helpful weather assistant.",
    tools=get_weather,
    middleware=[
        SecurityAgentMiddleware(),  # Applies to all runs
        TimingFunctionMiddleware(),  # Applies to all runs
    ],
) as agent:

    # This run uses agent-level middleware only
    result1 = await agent.run("What's the weather in Seattle?")

    # This run uses agent-level + run-level middleware
    result2 = await agent.run(
        "What's the weather in Portland?",
        middleware=[  # Run-level middleware (this run only)
            logging_chat_middleware,
        ]
    )

    # This run uses agent-level middleware only (no run-level)
    result3 = await agent.run("What's the weather in Vancouver?")

Diferencias clave:

• Nivel de agente: persistente en todas las ejecuciones, configuradas una vez al crear el agente
• Nivel de ejecución: solo se aplica a ejecuciones específicas, permite la personalización por solicitud.
• Orden de ejecución: middleware del agente (más externo) → Ejecución del middleware (más interno) → ejecución del agente

Finalización del middleware

El middleware puede finalizar la ejecución temprano estableciendo context.result y generando MiddlewareTermination. Esto resulta útil para las comprobaciones de seguridad, la limitación de velocidad o los errores de validación.

from agent_framework import AgentContext, AgentResponse, Message, MiddlewareTermination

async def blocking_middleware(
    context: AgentContext,
    call_next: Callable[[], Awaitable[None]],
) -> None:
    """Middleware that blocks execution based on conditions."""
    # Check for blocked content
    last_message = context.messages[-1] if context.messages else None
    if last_message and last_message.text:
        if "blocked" in last_message.text.lower():
            print("Request blocked by middleware")
            context.result = AgentResponse(
                messages=[Message(role="assistant", contents=["This request was blocked by middleware."])]
            )
            raise MiddlewareTermination(result=context.result)

    # If no issues, continue normally
    await call_next()

Qué terminación significa:

• Establecer context.result antes de generar MiddlewareTermination si desea devolver una respuesta personalizada
• La generación MiddlewareTermination detiene el resto de la cadena de middleware y omite la ruta de acceso de ejecución normal.
• Este patrón funciona para el middleware de agente, función y chat.

Invalidación de resultados de middleware

El middleware puede invalidar los resultados en escenarios que no son de streaming y streaming, lo que permite modificar o reemplazar completamente las respuestas del agente.

El tipo de resultado de context.result depende de si la invocación del agente es streaming o no de streaming:

• No streaming: context.result contiene un AgentResponse objeto con la respuesta completa.
• Streaming: context.result contiene un generador asincrónico que produce AgentResponseUpdate fragmentos.

Puede usar context.stream para diferenciar entre estos escenarios y controlar las invalidaciones de resultados correctamente.

async def weather_override_middleware(
    context: AgentContext,
    call_next: Callable[[], Awaitable[None]]
) -> None:
    """Middleware that overrides weather results for both streaming and non-streaming."""

    # Execute the original agent logic
    await call_next()

    # Override results if present
    if context.result is not None:
        custom_message_parts = [
            "Weather Override: ",
            "Perfect weather everywhere today! ",
            "22°C with gentle breezes. ",
            "Great day for outdoor activities!"
        ]

        if context.stream:
            # Streaming override
            async def override_stream() -> AsyncIterable[AgentResponseUpdate]:
                for chunk in custom_message_parts:
                    yield AgentResponseUpdate(contents=[Content.from_text(text=chunk)])

            context.result = override_stream()
        else:
            # Non-streaming override
            custom_message = "".join(custom_message_parts)
            context.result = AgentResponse(
                messages=[Message(role="assistant", contents=[custom_message])]
            )

Este enfoque de middleware permite implementar una sofisticada transformación de respuesta, filtrado de contenido, mejora de resultados y personalización de streaming, al tiempo que mantiene la lógica del agente limpia y centrada.

Ejemplos completos de middleware

Middleware basado en clases

# Copyright (c) Microsoft. All rights reserved.

import asyncio
import time
from collections.abc import Awaitable, Callable
from random import randint
from typing import Annotated

from agent_framework import (
    AgentContext,
    AgentMiddleware,
    AgentResponse,
    FunctionInvocationContext,
    FunctionMiddleware,
    Message,
    tool,
)
from agent_framework import Agent
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from pydantic import Field

"""
Class-based MiddlewareTypes Example

This sample demonstrates how to implement middleware using class-based approach by inheriting
from AgentMiddleware and FunctionMiddleware base classes. The example includes:

- SecurityAgentMiddleware: Checks for security violations in user queries and blocks requests
  containing sensitive information like passwords or secrets
- LoggingFunctionMiddleware: Logs function execution details including timing and parameters

This approach is useful when you need stateful middleware or complex logic that benefits
from object-oriented design patterns.
"""


# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
    location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
    """Get the weather for a given location."""
    conditions = ["sunny", "cloudy", "rainy", "stormy"]
    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."


class SecurityAgentMiddleware(AgentMiddleware):
    """Agent middleware that checks for security violations."""

    async def process(
        self,
        context: AgentContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        # Check for potential security violations in the query
        # Look at the last user message
        last_message = context.messages[-1] if context.messages else None
        if last_message and last_message.text:
            query = last_message.text
            if "password" in query.lower() or "secret" in query.lower():
                print("[SecurityAgentMiddleware] Security Warning: Detected sensitive information, blocking request.")
                # Override the result with warning message
                context.result = AgentResponse(
                    messages=[Message("assistant", ["Detected sensitive information, the request is blocked."])]
                )
                # Simply don't call call_next() to prevent execution
                return

        print("[SecurityAgentMiddleware] Security check passed.")
        await call_next()


class LoggingFunctionMiddleware(FunctionMiddleware):
    """Function middleware that logs function calls."""

    async def process(
        self,
        context: FunctionInvocationContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        function_name = context.function.name
        print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")

        start_time = time.time()

        await call_next()

        end_time = time.time()
        duration = end_time - start_time

        print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.5f}s.")


async def main() -> None:
    """Example demonstrating class-based middleware."""
    print("=== Class-based MiddlewareTypes Example ===")

    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
    # authentication option.
    async with (
        AzureCliCredential() as credential,
        Agent(
            client=FoundryChatClient(credential=credential),
            name="WeatherAgent",
            instructions="You are a helpful weather assistant.",
            tools=get_weather,
            middleware=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()],
        ) as agent,
    ):
        # Test with normal query
        print("\n--- Normal Query ---")
        query = "What's the weather like in Seattle?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}\n")

        # Test with security-related query
        print("--- Security Test ---")
        query = "What's the password for the weather service?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}\n")


if __name__ == "__main__":
    asyncio.run(main())

Middleware basado en funciones

# Copyright (c) Microsoft. All rights reserved.

import asyncio
import time
from collections.abc import Awaitable, Callable
from random import randint
from typing import Annotated

from agent_framework import (
    AgentContext,
    AgentMiddleware,
    AgentResponse,
    FunctionInvocationContext,
    FunctionMiddleware,
    Message,
    tool,
)
from agent_framework import Agent
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from pydantic import Field

"""
Class-based MiddlewareTypes Example

This sample demonstrates how to implement middleware using class-based approach by inheriting
from AgentMiddleware and FunctionMiddleware base classes. The example includes:

- SecurityAgentMiddleware: Checks for security violations in user queries and blocks requests
  containing sensitive information like passwords or secrets
- LoggingFunctionMiddleware: Logs function execution details including timing and parameters

This approach is useful when you need stateful middleware or complex logic that benefits
from object-oriented design patterns.
"""


# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
    location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
    """Get the weather for a given location."""
    conditions = ["sunny", "cloudy", "rainy", "stormy"]
    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."


class SecurityAgentMiddleware(AgentMiddleware):
    """Agent middleware that checks for security violations."""

    async def process(
        self,
        context: AgentContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        # Check for potential security violations in the query
        # Look at the last user message
        last_message = context.messages[-1] if context.messages else None
        if last_message and last_message.text:
            query = last_message.text
            if "password" in query.lower() or "secret" in query.lower():
                print("[SecurityAgentMiddleware] Security Warning: Detected sensitive information, blocking request.")
                # Override the result with warning message
                context.result = AgentResponse(
                    messages=[Message("assistant", ["Detected sensitive information, the request is blocked."])]
                )
                # Simply don't call call_next() to prevent execution
                return

        print("[SecurityAgentMiddleware] Security check passed.")
        await call_next()


class LoggingFunctionMiddleware(FunctionMiddleware):
    """Function middleware that logs function calls."""

    async def process(
        self,
        context: FunctionInvocationContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        function_name = context.function.name
        print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")

        start_time = time.time()

        await call_next()

        end_time = time.time()
        duration = end_time - start_time

        print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.5f}s.")


async def main() -> None:
    """Example demonstrating class-based middleware."""
    print("=== Class-based MiddlewareTypes Example ===")

    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
    # authentication option.
    async with (
        AzureCliCredential() as credential,
        Agent(
            client=FoundryChatClient(credential=credential),
            name="WeatherAgent",
            instructions="You are a helpful weather assistant.",
            tools=get_weather,
            middleware=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()],
        ) as agent,
    ):
        # Test with normal query
        print("\n--- Normal Query ---")
        query = "What's the weather like in Seattle?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}\n")

        # Test with security-related query
        print("--- Security Test ---")
        query = "What's the password for the weather service?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}\n")


if __name__ == "__main__":
    asyncio.run(main())

Middleware basado en decorador

# Copyright (c) Microsoft. All rights reserved.

import asyncio
import time
from collections.abc import Awaitable, Callable
from random import randint
from typing import Annotated

from agent_framework import (
    AgentContext,
    AgentMiddleware,
    AgentResponse,
    FunctionInvocationContext,
    FunctionMiddleware,
    Message,
    tool,
)
from agent_framework import Agent
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from pydantic import Field

"""
Class-based MiddlewareTypes Example

This sample demonstrates how to implement middleware using class-based approach by inheriting
from AgentMiddleware and FunctionMiddleware base classes. The example includes:

- SecurityAgentMiddleware: Checks for security violations in user queries and blocks requests
  containing sensitive information like passwords or secrets
- LoggingFunctionMiddleware: Logs function execution details including timing and parameters

This approach is useful when you need stateful middleware or complex logic that benefits
from object-oriented design patterns.
"""


# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
    location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
    """Get the weather for a given location."""
    conditions = ["sunny", "cloudy", "rainy", "stormy"]
    return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."


class SecurityAgentMiddleware(AgentMiddleware):
    """Agent middleware that checks for security violations."""

    async def process(
        self,
        context: AgentContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        # Check for potential security violations in the query
        # Look at the last user message
        last_message = context.messages[-1] if context.messages else None
        if last_message and last_message.text:
            query = last_message.text
            if "password" in query.lower() or "secret" in query.lower():
                print("[SecurityAgentMiddleware] Security Warning: Detected sensitive information, blocking request.")
                # Override the result with warning message
                context.result = AgentResponse(
                    messages=[Message("assistant", ["Detected sensitive information, the request is blocked."])]
                )
                # Simply don't call call_next() to prevent execution
                return

        print("[SecurityAgentMiddleware] Security check passed.")
        await call_next()


class LoggingFunctionMiddleware(FunctionMiddleware):
    """Function middleware that logs function calls."""

    async def process(
        self,
        context: FunctionInvocationContext,
        call_next: Callable[[], Awaitable[None]],
    ) -> None:
        function_name = context.function.name
        print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")

        start_time = time.time()

        await call_next()

        end_time = time.time()
        duration = end_time - start_time

        print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.5f}s.")


async def main() -> None:
    """Example demonstrating class-based middleware."""
    print("=== Class-based MiddlewareTypes Example ===")

    # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
    # authentication option.
    async with (
        AzureCliCredential() as credential,
        Agent(
            client=FoundryChatClient(credential=credential),
            name="WeatherAgent",
            instructions="You are a helpful weather assistant.",
            tools=get_weather,
            middleware=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()],
        ) as agent,
    ):
        # Test with normal query
        print("\n--- Normal Query ---")
        query = "What's the weather like in Seattle?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}\n")

        # Test with security-related query
        print("--- Security Test ---")
        query = "What's the password for the weather service?"
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}\n")


if __name__ == "__main__":
    asyncio.run(main())