Note
Access to this page requires authorization. You can try signing in or changing directories.
Access to this page requires authorization. You can try changing directories.
Middleware in Agent Framework provides a powerful way to intercept, modify, and enhance agent interactions at various stages of execution. You can use middleware to implement cross-cutting concerns such as logging, security validation, error handling, and result transformation without modifying your core agent or function logic.
Agent Framework can be customized using three different types of middleware:
- Agent Run middleware: Allows interception of all agent runs, so that input and output can be inspected and/or modified as needed.
- Function calling middleware: Allows interception of all function calls executed by the agent, so that input and output can be inspected and modified as needed.
- IChatClient middleware: Allows interception of calls to an
IChatClientimplementation, where an agent is usingIChatClientfor inference calls, for example, when usingChatClientAgent.
All the types of middleware are implemented via a function callback, and when multiple middleware instances of the same type are registered, they form a chain,
where each middleware instance is expected to call the next in the chain, via a provided next Func.
Agent run and function calling middleware types can be registered on an agent, by using the agent builder with an existing agent object.
var middlewareEnabledAgent = originalAgent
.AsBuilder()
.Use(runFunc: CustomAgentRunMiddleware, runStreamingFunc: CustomAgentRunStreamingMiddleware)
.Use(CustomFunctionCallingMiddleware)
.Build();
Important
Ideally both runFunc and runStreamingFunc should be provided. When providing just the non-streaming middleware, the agent will use it for both streaming and non-streaming invocations. Streaming will only run in non-streaming mode to suffice the middleware expectations.
Note
There's an additional overload, Use(sharedFunc: ...), that allows you to provide the same middleware for non-streaming and streaming without blocking the streaming. However, the shared middleware won't be able to intercept or override the output. This overload should be used for scenarios where you only need to inspect or modify the input before it reaches the agent.
IChatClient middleware can be registered on an IChatClient before it is used with a ChatClientAgent, by using the chat client builder pattern.
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.");
Warning
DefaultAzureCredential is convenient for development but requires careful consideration in production. In production, consider using a specific credential (e.g., ManagedIdentityCredential) to avoid latency issues, unintended credential probing, and potential security risks from fallback mechanisms.
IChatClient middleware can also be registered using a factory method when constructing
an agent via one of the helper methods on SDK clients.
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());
Agent Run Middleware
Here is an example of agent run middleware, that can inspect and/or modify the input and output from the agent run.
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;
}
Agent Run Streaming Middleware
Here is an example of agent run streaming middleware, that can inspect and/or modify the input and output from the agent streaming run.
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);
}
Function calling middleware
Note
Function calling middleware is currently only supported with an AIAgent that uses FunctionInvokingChatClient, for example, ChatClientAgent.
Here is an example of function calling middleware, that can inspect and/or modify the function being called, and the result from the function call.
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;
}
It is possible to terminate the function call loop with function calling middleware by setting the provided FunctionInvocationContext.Terminate to true.
This will prevent the function calling loop from issuing a request to the inference service containing the function call results after function invocation.
If there were more than one function available for invocation during this iteration, it might also prevent any remaining functions from being executed.
Warning
Terminating the function call loop might result in your chat history being left in an inconsistent state, for example, containing function call content with no function result content. This might result in the chat history being unusable for further runs.
IChatClient middleware
Here is an example of chat client middleware, that can inspect and/or modify the input and output for the request to the inference service that the chat client provides.
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;
}
Tip
See the .NET samples for complete runnable examples.
Note
For more information about IChatClient middleware, see Custom IChatClient middleware.
Agent Framework can be customized using three different types of middleware:
- Agent middleware: Intercepts agent run execution, allowing you to inspect and modify inputs, outputs, and control flow.
- Function middleware: Intercepts function (tool) calls made during agent execution, enabling input validation, result transformation, and execution control.
- Chat middleware: Intercepts the underlying chat requests sent to AI models, providing access to the raw messages, options, and responses.
All types support both function-based and class-based implementations. When multiple middleware of the same type are registered, they form a chain where each calls the next callable to continue processing.
Agent Middleware
Agent middleware intercepts and modifies agent run execution. It uses the AgentContext which contains:
agent: The agent being invokedmessages: List of chat messages in the conversationis_streaming: Boolean indicating if the response is streamingmetadata: Dictionary for storing additional data between middlewareresult: The agent's response (can be modified)terminate: Flag to stop further processingkwargs: Additional keyword arguments passed to the agent run method
The next callable continues the middleware chain or executes the agent if it's the last middleware.
Function-based
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")
Class-based
Class-based agent middleware uses a process method that has the same signature and behavior as function-based middleware.
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")
Function Middleware
Function middleware intercepts function calls within agents. It uses the FunctionInvocationContext which contains:
function: The function being invokedarguments: The validated arguments for the functionmetadata: Dictionary for storing additional data between middlewareresult: The function's return value (can be modified)terminate: Flag to stop further processingkwargs: Additional keyword arguments passed to the chat method that invoked this function
The next callable continues to the next middleware or executes the actual function.
Function-based
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")
Class-based
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
Chat middleware intercepts chat requests sent to AI models. It uses the ChatContext which contains:
chat_client: The chat client being invokedmessages: List of messages being sent to the AI serviceoptions: The options for the chat requestis_streaming: Boolean indicating if this is a streaming invocationmetadata: Dictionary for storing additional data between middlewareresult: The chat response from the AI (can be modified)terminate: Flag to stop further processingkwargs: Additional keyword arguments passed to the chat client
The next callable continues to the next middleware or sends the request to the AI service.
Function-based
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")
Class-based
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")
Middleware Decorators
Decorators provide explicit middleware type declaration without requiring type annotations. They're helpful when you don't use type annotations or want to prevent type mismatches:
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")
Middleware Registration
Middleware can be registered at two levels with different scopes and behaviors.
Agent-Level vs Run-Level Middleware
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?")
Key Differences:
- Agent-level: Persistent across all runs, configured once when creating the agent
- Run-level: Applied only to specific runs, allows per-request customization
- Execution Order: Agent middleware (outermost) → Run middleware (innermost) → Agent execution
Middleware Termination
Middleware can terminate execution early using context.terminate. This is useful for security checks, rate limiting, or validation failures.
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)
What termination means:
- Setting
context.terminate = Truesignals that processing should stop - You can provide a custom result before terminating to give users feedback
- The agent execution is completely skipped when middleware terminates
Middleware Result Override
Middleware can override results in both non-streaming and streaming scenarios, allowing you to modify or completely replace agent responses.
The result type in context.result depends on whether the agent invocation is streaming or non-streaming:
- Non-streaming:
context.resultcontains anAgentResponsewith the complete response - Streaming:
context.resultcontains an async generator that yieldsAgentResponseUpdatechunks
You can use context.is_streaming to differentiate between these scenarios and handle result overrides appropriately.
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])]
)
This middleware approach allows you to implement sophisticated response transformation, content filtering, result enhancement, and streaming customization while keeping your agent logic clean and focused.
Complete middleware examples
Class-based middleware
# 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())
Function-based middleware
# 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())
Decorator-based middleware
# 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())