Nota:
El acceso a esta página requiere autorización. Puede intentar iniciar sesión o cambiar directorios.
El acceso a esta página requiere autorización. Puede intentar cambiar los directorios.
El middleware de Agent Framework proporciona una manera eficaz de interceptar, modificar y mejorar las interacciones del agente en varias fases de ejecución. Puede usar middleware para implementar problemas transversales como el registro, la validación de seguridad, el control de errores y la transformación de resultados sin modificar el agente principal ni la lógica de la función.
Agent Framework se puede personalizar mediante tres tipos diferentes de middleware:
- 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.
- 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.
-
IChatClient middleware: permite la interceptación de llamadas a una
IChatClientimplementación, donde un agente usaIChatClientpara llamadas de inferencia, por ejemplo, al usarChatClientAgent.
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();
Importante
Lo ideal es que se proporcionen y runFuncrunStreamingFunc . Al proporcionar solo el middleware que no es de streaming, el agente lo usará tanto para invocaciones de streaming como para no streaming. El streaming solo se ejecutará en modo que no sea de streaming para bastar con las expectativas de middleware.
Nota:
Hay una sobrecarga adicional, Use(sharedFunc: ...), que permite proporcionar el mismo middleware para no streaming y streaming sin bloquear el streaming. Sin embargo, el middleware compartido no podrá interceptar ni invalidar la salida. Esta sobrecarga debe usarse en escenarios en los que solo es necesario inspeccionar o modificar la entrada antes de que llegue al agente.
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 AzureOpenAIClient(new Uri("https://<myresource>.openai.azure.com"), new DefaultAzureCredential())
.GetChatClient(deploymentName)
.AsIChatClient();
var middlewareEnabledChatClient = chatClient
.AsBuilder()
.Use(getResponseFunc: CustomChatClientMiddleware, getStreamingResponseFunc: null)
.Build();
var agent = new ChatClientAgent(middlewareEnabledChatClient, instructions: "You are a helpful assistant.");
Advertencia
DefaultAzureCredential es conveniente para el desarrollo, pero requiere una consideración cuidadosa en producción. En producción, considere la posibilidad de usar una credencial específica (por ejemplo, ManagedIdentityCredential) para evitar problemas de latencia, sondeos de credenciales no deseados y posibles riesgos de seguridad de los mecanismos de reserva.
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 AzureOpenAIClient(new Uri(endpoint), new DefaultAzureCredential())
.GetChatClient(deploymentName)
.AsAIAgent("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
Nota:
Actualmente, el middleware de llamada a funciones solo se admite con un AIAgent que usa FunctionInvokingChatClient, por ejemplo, ChatClientAgent.
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.
Advertencia
La terminación del bucle de llamada de función puede dar lugar a que el historial de chat se quede en un estado incoherente, por ejemplo, que contenga contenido de llamada de función sin contenido de resultado de función. Esto puede dar lugar a que el historial de chat no se pueda usar para más ejecuciones.
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;
}
Sugerencia
Consulte los ejemplos de .NET para obtener ejemplos completos de ejecución.
Nota:
Para obtener más información sobre IChatClient el middleware, consulte middleware personalizado de IChatClient.
Agent Framework se puede personalizar mediante tres tipos diferentes de middleware:
- 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.
- 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.
- 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 al que se puede next seguir procesando.
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 -
is_streaming: booleano que indica si la respuesta es streaming. -
metadata: diccionario para almacenar datos adicionales entre middleware -
result: la respuesta del agente (se puede modificar) -
terminate: marca para detener el procesamiento adicional -
kwargs: argumentos de palabra clave adicionales pasados al método run del agente
El next invocable continúa la cadena de middleware o ejecuta el agente si es el último middleware.
Basado en funciones
async def logging_agent_middleware(
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
) -> None:
"""Agent middleware that logs execution timing."""
# Pre-processing: Log before agent execution
print("[Agent] Starting execution")
# Continue to next middleware or agent execution
await next(context)
# Post-processing: Log after agent execution
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,
next: Callable[[AgentContext], Awaitable[None]],
) -> None:
print("[Agent Class] Starting execution")
await next(context)
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 -
metadata: diccionario para almacenar datos adicionales entre middleware -
result: el valor devuelto de la función (se puede modificar) -
terminate: marca para detener el procesamiento adicional -
kwargs: argumentos de palabra clave adicionales pasados al método de chat que invocó esta función
El next invocable continúa con el siguiente middleware o ejecuta la función real.
Basado en funciones
async def logging_function_middleware(
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Function middleware that logs function execution."""
# Pre-processing: Log before function execution
print(f"[Function] Calling {context.function.name}")
# Continue to next middleware or function execution
await next(context)
# Post-processing: Log after function execution
print(f"[Function] {context.function.name} completed")
Basado en clases
from agent_framework import FunctionMiddleware, FunctionInvocationContext
class LoggingFunctionMiddleware(FunctionMiddleware):
"""Function middleware that logs function execution."""
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
print(f"[Function Class] Calling {context.function.name}")
await next(context)
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 -
is_streaming: 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) -
terminate: marca para detener el procesamiento adicional -
kwargs: argumentos de palabra clave adicionales pasados al cliente de chat
El next invocable 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,
next: Callable[[ChatContext], 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 next(context)
# 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,
next: Callable[[ChatContext], Awaitable[None]],
) -> None:
print(f"[Chat Class] Sending {len(context.messages)} messages to AI")
await next(context)
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, next):
print("Before agent execution")
await next(context)
print("After agent execution")
@function_middleware
async def simple_function_middleware(context, next):
print(f"Calling function: {context.function.name}")
await next(context)
print("Function call completed")
@chat_middleware
async def simple_chat_middleware(context, next):
print(f"Processing {len(context.messages)} chat messages")
await next(context)
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.azure import AzureAIAgentClient
from azure.identity.aio import AzureCliCredential
# Agent-level middleware: Applied to ALL runs of the agent
async with AzureAIAgentClient(async_credential=credential).as_agent(
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 al principio mediante context.terminate. Esto resulta útil para las comprobaciones de seguridad, la limitación de velocidad o los errores de validación.
async def blocking_middleware(
context: AgentContext,
next: Callable[[AgentContext], 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.terminate = True
return
# If no issues, continue normally
await next(context)
Qué terminación significa:
- Establecer
context.terminate = Trueseñales de que el procesamiento debe detenerse - Puede proporcionar un resultado personalizado antes de terminar para proporcionar comentarios a los usuarios.
- La ejecución del agente se omite completamente cuando finaliza el middleware.
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.resultcontiene unAgentResponseobjeto con la respuesta completa. -
Streaming:
context.resultcontiene un generador asincrónico que produceAgentResponseUpdatefragmentos.
Puede usar context.is_streaming para diferenciar entre estos escenarios y controlar las invalidaciones de resultados correctamente.
async def weather_override_middleware(
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]]
) -> None:
"""Middleware that overrides weather results for both streaming and non-streaming."""
# Execute the original agent logic
await next(context)
# 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.is_streaming:
# 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.azure import AzureAIAgentClient
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,
AzureAIAgentClient(credential=credential).as_agent(
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.azure import AzureAIAgentClient
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,
AzureAIAgentClient(credential=credential).as_agent(
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.azure import AzureAIAgentClient
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,
AzureAIAgentClient(credential=credential).as_agent(
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())