test
Some checks failed
CodeQL / Analyze (csharp) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
Python - Merge - Tests / paths-filter (push) Has been cancelled
Python - Merge - Tests / Python Tests - Core (integration, ubuntu-latest, 3.10) (push) Has been cancelled
Python - Merge - Tests / Python Tests - Azure AI (integration, ubuntu-latest, 3.10) (push) Has been cancelled
Python - Merge - Tests / python-integration-tests-check (push) Has been cancelled
Python - Lab Tests / paths-filter (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.10) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.11) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.12) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.13) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.14) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.10) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.11) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.12) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.13) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.14) (push) Has been cancelled
Check .md links / markdown-link-check (push) Has been cancelled
Some checks failed
CodeQL / Analyze (csharp) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
Python - Merge - Tests / paths-filter (push) Has been cancelled
Python - Merge - Tests / Python Tests - Core (integration, ubuntu-latest, 3.10) (push) Has been cancelled
Python - Merge - Tests / Python Tests - Azure AI (integration, ubuntu-latest, 3.10) (push) Has been cancelled
Python - Merge - Tests / python-integration-tests-check (push) Has been cancelled
Python - Lab Tests / paths-filter (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.10) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.11) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.12) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.13) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (ubuntu-latest, 3.14) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.10) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.11) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.12) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.13) (push) Has been cancelled
Python - Lab Tests / Python Lab Tests (windows-latest, 3.14) (push) Has been cancelled
Check .md links / markdown-link-check (push) Has been cancelled
This commit is contained in:
46
python/samples/getting_started/middleware/README.md
Normal file
46
python/samples/getting_started/middleware/README.md
Normal file
@@ -0,0 +1,46 @@
|
||||
# Middleware Examples
|
||||
|
||||
This folder contains examples demonstrating various middleware patterns with the Agent Framework. Middleware allows you to intercept and modify behavior at different execution stages, including agent runs, function calls, and chat interactions.
|
||||
|
||||
## Examples
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| [`function_based_middleware.py`](function_based_middleware.py) | Demonstrates how to implement middleware using simple async functions instead of classes. Shows security validation, logging, and performance monitoring middleware. Function-based middleware is ideal for simple, stateless operations and provides a lightweight approach. |
|
||||
| [`class_based_middleware.py`](class_based_middleware.py) | Shows how to implement middleware using class-based approach by inheriting from `AgentMiddleware` and `FunctionMiddleware` base classes. Includes security checks for sensitive information and detailed function execution logging with timing. |
|
||||
| [`decorator_middleware.py`](decorator_middleware.py) | Demonstrates how to use `@agent_middleware` and `@function_middleware` decorators to explicitly mark middleware functions without requiring type annotations. Shows different middleware detection scenarios and explicit decorator usage. |
|
||||
| [`middleware_termination.py`](middleware_termination.py) | Shows how middleware can terminate execution using the `context.terminate` flag. Includes examples of pre-termination (prevents agent processing) and post-termination (allows processing but stops further execution). Useful for security checks, rate limiting, or early exit conditions. |
|
||||
| [`exception_handling_with_middleware.py`](exception_handling_with_middleware.py) | Demonstrates how to use middleware for centralized exception handling in function calls. Shows how to catch exceptions from functions, provide graceful error responses, and override function results when errors occur to provide user-friendly messages. |
|
||||
| [`override_result_with_middleware.py`](override_result_with_middleware.py) | Shows how to use middleware to intercept and modify function results after execution, supporting both regular and streaming agent responses. Demonstrates result filtering, formatting, enhancement, and custom streaming response generation. |
|
||||
| [`shared_state_middleware.py`](shared_state_middleware.py) | Demonstrates how to implement function-based middleware within a class to share state between multiple middleware functions. Shows how middleware can work together by sharing state, including call counting and result enhancement. |
|
||||
| [`thread_behavior_middleware.py`](thread_behavior_middleware.py) | Demonstrates how middleware can access and track thread state across multiple agent runs. Shows how `AgentRunContext.thread` behaves differently before and after the `next()` call, how conversation history accumulates in threads, and timing of thread message updates. Essential for understanding conversation flow in middleware. |
|
||||
| [`agent_and_run_level_middleware.py`](agent_and_run_level_middleware.py) | Explains the difference between agent-level middleware (applied to ALL runs of the agent) and run-level middleware (applied to specific runs only). Shows security validation, performance monitoring, and context-specific middleware patterns. |
|
||||
| [`chat_middleware.py`](chat_middleware.py) | Demonstrates how to use chat middleware to observe and override inputs sent to AI models. Shows how to intercept chat requests, log and modify input messages, and override entire responses before they reach the underlying AI service. |
|
||||
|
||||
## Key Concepts
|
||||
|
||||
### Middleware Types
|
||||
|
||||
- **Agent Middleware**: Intercepts agent run execution, allowing you to modify requests and responses
|
||||
- **Function Middleware**: Intercepts function calls within agents, enabling logging, validation, and result modification
|
||||
- **Chat Middleware**: Intercepts chat requests sent to AI models, allowing input/output transformation
|
||||
|
||||
### Implementation Approaches
|
||||
|
||||
- **Function-based**: Simple async functions for lightweight, stateless operations
|
||||
- **Class-based**: Inherit from base middleware classes for complex, stateful operations
|
||||
- **Decorator-based**: Use decorators for explicit middleware marking
|
||||
|
||||
### Common Use Cases
|
||||
|
||||
- **Security**: Validate requests, block sensitive information, implement access controls
|
||||
- **Logging**: Track execution timing, log parameters and results, monitor performance
|
||||
- **Error Handling**: Catch exceptions, provide graceful fallbacks, implement retry logic
|
||||
- **Result Transformation**: Filter, format, or enhance function outputs
|
||||
- **State Management**: Share data between middleware functions, maintain execution context
|
||||
|
||||
### Execution Control
|
||||
|
||||
- **Termination**: Use `context.terminate` to stop execution early
|
||||
- **Result Override**: Modify or replace function/agent results
|
||||
- **Streaming Support**: Handle both regular and streaming responses
|
||||
@@ -0,0 +1,271 @@
|
||||
# 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 (
|
||||
AgentMiddleware,
|
||||
AgentResponse,
|
||||
AgentRunContext,
|
||||
FunctionInvocationContext,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Agent-Level and Run-Level Middleware Example
|
||||
|
||||
This sample demonstrates the difference between agent-level and run-level middleware:
|
||||
|
||||
- Agent-level middleware: Applied to ALL runs of the agent (persistent across runs)
|
||||
- Run-level middleware: Applied to specific runs only (isolated per run)
|
||||
|
||||
The example shows:
|
||||
1. Agent-level security middleware that validates all requests
|
||||
2. Agent-level performance monitoring across all runs
|
||||
3. Run-level context middleware for specific use cases (high priority, debugging)
|
||||
4. Run-level caching middleware for expensive operations
|
||||
|
||||
Execution order: Agent middleware (outermost) -> Run middleware (innermost) -> Agent execution
|
||||
"""
|
||||
|
||||
|
||||
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."
|
||||
|
||||
|
||||
# Agent-level middleware (applied to ALL runs)
|
||||
class SecurityAgentMiddleware(AgentMiddleware):
|
||||
"""Agent-level security middleware that validates all requests."""
|
||||
|
||||
async def process(self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]) -> None:
|
||||
print("[SecurityMiddleware] Checking security for all requests...")
|
||||
|
||||
# Check for security violations in 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.lower()
|
||||
if any(word in query for word in ["password", "secret", "credentials"]):
|
||||
print("[SecurityMiddleware] Security violation detected! Blocking request.")
|
||||
return # Don't call next() to prevent execution
|
||||
|
||||
print("[SecurityMiddleware] Security check passed.")
|
||||
context.metadata["security_validated"] = True
|
||||
await next(context)
|
||||
|
||||
|
||||
async def performance_monitor_middleware(
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Agent-level performance monitoring for all runs."""
|
||||
print("[PerformanceMonitor] Starting performance monitoring...")
|
||||
start_time = time.time()
|
||||
|
||||
await next(context)
|
||||
|
||||
end_time = time.time()
|
||||
duration = end_time - start_time
|
||||
print(f"[PerformanceMonitor] Total execution time: {duration:.3f}s")
|
||||
context.metadata["execution_time"] = duration
|
||||
|
||||
|
||||
# Run-level middleware (applied to specific runs only)
|
||||
class HighPriorityMiddleware(AgentMiddleware):
|
||||
"""Run-level middleware for high priority requests."""
|
||||
|
||||
async def process(self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]) -> None:
|
||||
print("[HighPriority] Processing high priority request with expedited handling...")
|
||||
|
||||
# Read metadata set by agent-level middleware
|
||||
if context.metadata.get("security_validated"):
|
||||
print("[HighPriority] Security validation confirmed from agent middleware")
|
||||
|
||||
# Set high priority flag
|
||||
context.metadata["priority"] = "high"
|
||||
context.metadata["expedited"] = True
|
||||
|
||||
await next(context)
|
||||
print("[HighPriority] High priority processing completed")
|
||||
|
||||
|
||||
async def debugging_middleware(
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Run-level debugging middleware for troubleshooting specific runs."""
|
||||
print("[Debug] Debug mode enabled for this run")
|
||||
print(f"[Debug] Messages count: {len(context.messages)}")
|
||||
print(f"[Debug] Is streaming: {context.is_streaming}")
|
||||
|
||||
# Log existing metadata from agent middleware
|
||||
if context.metadata:
|
||||
print(f"[Debug] Existing metadata: {context.metadata}")
|
||||
|
||||
context.metadata["debug_enabled"] = True
|
||||
|
||||
await next(context)
|
||||
|
||||
print("[Debug] Debug information collected")
|
||||
|
||||
|
||||
class CachingMiddleware(AgentMiddleware):
|
||||
"""Run-level caching middleware for expensive operations."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.cache: dict[str, AgentResponse] = {}
|
||||
|
||||
async def process(self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]) -> None:
|
||||
# Create a simple cache key from the last message
|
||||
last_message = context.messages[-1] if context.messages else None
|
||||
cache_key: str = last_message.text if last_message and last_message.text else "no_message"
|
||||
|
||||
if cache_key in self.cache:
|
||||
print(f"[Cache] Cache HIT for: '{cache_key[:30]}...'")
|
||||
context.result = self.cache[cache_key] # type: ignore
|
||||
return # Don't call next(), return cached result
|
||||
|
||||
print(f"[Cache] Cache MISS for: '{cache_key[:30]}...'")
|
||||
context.metadata["cache_key"] = cache_key
|
||||
|
||||
await next(context)
|
||||
|
||||
# Cache the result if we have one
|
||||
if context.result:
|
||||
self.cache[cache_key] = context.result # type: ignore
|
||||
print("[Cache] Result cached for future use")
|
||||
|
||||
|
||||
async def function_logging_middleware(
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Function middleware that logs all function calls."""
|
||||
function_name = context.function.name
|
||||
args = context.arguments
|
||||
print(f"[FunctionLog] Calling function: {function_name} with args: {args}")
|
||||
|
||||
await next(context)
|
||||
|
||||
print(f"[FunctionLog] Function {function_name} completed")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating agent-level and run-level middleware."""
|
||||
print("=== Agent-Level and Run-Level Middleware Example ===\n")
|
||||
|
||||
# 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,
|
||||
# Agent-level middleware: applied to ALL runs
|
||||
middleware=[
|
||||
SecurityAgentMiddleware(),
|
||||
performance_monitor_middleware,
|
||||
function_logging_middleware,
|
||||
],
|
||||
) as agent,
|
||||
):
|
||||
print("Agent created with agent-level middleware:")
|
||||
print(" - SecurityMiddleware (blocks sensitive requests)")
|
||||
print(" - PerformanceMonitor (tracks execution time)")
|
||||
print(" - FunctionLogging (logs all function calls)")
|
||||
print()
|
||||
|
||||
# Run 1: Normal query with no run-level middleware
|
||||
print("=" * 60)
|
||||
print("RUN 1: Normal query (agent-level middleware only)")
|
||||
print("=" * 60)
|
||||
query = "What's the weather like in Paris?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
print()
|
||||
|
||||
# Run 2: High priority request with run-level middleware
|
||||
print("=" * 60)
|
||||
print("RUN 2: High priority request (agent + run-level middleware)")
|
||||
print("=" * 60)
|
||||
query = "What's the weather in Tokyo? This is urgent!"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(
|
||||
query,
|
||||
middleware=[HighPriorityMiddleware()], # Run-level middleware
|
||||
)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
print()
|
||||
|
||||
# Run 3: Debug mode with run-level debugging middleware
|
||||
print("=" * 60)
|
||||
print("RUN 3: Debug mode (agent + run-level debugging)")
|
||||
print("=" * 60)
|
||||
query = "What's the weather in London?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(
|
||||
query,
|
||||
middleware=[debugging_middleware], # Run-level middleware
|
||||
)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
print()
|
||||
|
||||
# Run 4: Multiple run-level middleware
|
||||
print("=" * 60)
|
||||
print("RUN 4: Multiple run-level middleware (caching + debug)")
|
||||
print("=" * 60)
|
||||
caching = CachingMiddleware()
|
||||
query = "What's the weather in New York?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(
|
||||
query,
|
||||
middleware=[caching, debugging_middleware], # Multiple run-level middleware
|
||||
)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
print()
|
||||
|
||||
# Run 5: Test cache hit with same query
|
||||
print("=" * 60)
|
||||
print("RUN 5: Test cache hit (same query as Run 4)")
|
||||
print("=" * 60)
|
||||
print(f"User: {query}") # Same query as Run 4
|
||||
result = await agent.run(
|
||||
query,
|
||||
middleware=[caching], # Same caching middleware instance
|
||||
)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
print()
|
||||
|
||||
# Run 6: Security violation test
|
||||
print("=" * 60)
|
||||
print("RUN 6: Security test (should be blocked by agent middleware)")
|
||||
print("=" * 60)
|
||||
query = "What's the secret weather password for Berlin?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'Request was blocked by security middleware'}")
|
||||
print()
|
||||
|
||||
# Run 7: Normal query again (no run-level middleware interference)
|
||||
print("=" * 60)
|
||||
print("RUN 7: Normal query again (agent-level middleware only)")
|
||||
print("=" * 60)
|
||||
query = "What's the weather in Sydney?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
245
python/samples/getting_started/middleware/chat_middleware.py
Normal file
245
python/samples/getting_started/middleware/chat_middleware.py
Normal file
@@ -0,0 +1,245 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from random import randint
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import (
|
||||
ChatContext,
|
||||
ChatMessage,
|
||||
ChatMiddleware,
|
||||
ChatResponse,
|
||||
Role,
|
||||
chat_middleware,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Chat Middleware Example
|
||||
|
||||
This sample demonstrates how to use chat middleware to observe and override
|
||||
inputs sent to AI models. Chat middleware intercepts chat requests before they reach
|
||||
the underlying AI service, allowing you to:
|
||||
|
||||
1. Observe and log input messages
|
||||
2. Modify input messages before sending to AI
|
||||
3. Override the entire response
|
||||
|
||||
The example covers:
|
||||
- Class-based chat middleware inheriting from ChatMiddleware
|
||||
- Function-based chat middleware with @chat_middleware decorator
|
||||
- Middleware registration at agent level (applies to all runs)
|
||||
- Middleware registration at run level (applies to specific run only)
|
||||
"""
|
||||
|
||||
|
||||
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 InputObserverMiddleware(ChatMiddleware):
|
||||
"""Class-based middleware that observes and modifies input messages."""
|
||||
|
||||
def __init__(self, replacement: str | None = None):
|
||||
"""Initialize with a replacement for user messages."""
|
||||
self.replacement = replacement
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: ChatContext,
|
||||
next: Callable[[ChatContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Observe and modify input messages before they are sent to AI."""
|
||||
print("[InputObserverMiddleware] Observing input messages:")
|
||||
|
||||
for i, message in enumerate(context.messages):
|
||||
content = message.text if message.text else str(message.contents)
|
||||
print(f" Message {i + 1} ({message.role.value}): {content}")
|
||||
|
||||
print(f"[InputObserverMiddleware] Total messages: {len(context.messages)}")
|
||||
|
||||
# Modify user messages by creating new messages with enhanced text
|
||||
modified_messages: list[ChatMessage] = []
|
||||
modified_count = 0
|
||||
|
||||
for message in context.messages:
|
||||
if message.role == Role.USER and message.text:
|
||||
original_text = message.text
|
||||
updated_text = original_text
|
||||
|
||||
if self.replacement:
|
||||
updated_text = self.replacement
|
||||
print(f"[InputObserverMiddleware] Updated: '{original_text}' -> '{updated_text}'")
|
||||
|
||||
modified_message = ChatMessage(role=message.role, text=updated_text)
|
||||
modified_messages.append(modified_message)
|
||||
modified_count += 1
|
||||
else:
|
||||
modified_messages.append(message)
|
||||
|
||||
# Replace messages in context
|
||||
context.messages[:] = modified_messages
|
||||
|
||||
# Continue to next middleware or AI execution
|
||||
await next(context)
|
||||
|
||||
# Observe that processing is complete
|
||||
print("[InputObserverMiddleware] Processing completed")
|
||||
|
||||
|
||||
@chat_middleware
|
||||
async def security_and_override_middleware(
|
||||
context: ChatContext,
|
||||
next: Callable[[ChatContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Function-based middleware that implements security filtering and response override."""
|
||||
print("[SecurityMiddleware] Processing input...")
|
||||
|
||||
# Security check - block sensitive information
|
||||
blocked_terms = ["password", "secret", "api_key", "token"]
|
||||
|
||||
for message in context.messages:
|
||||
if message.text:
|
||||
message_lower = message.text.lower()
|
||||
for term in blocked_terms:
|
||||
if term in message_lower:
|
||||
print(f"[SecurityMiddleware] BLOCKED: Found '{term}' in message")
|
||||
|
||||
# Override the response instead of calling AI
|
||||
context.result = ChatResponse(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
text="I cannot process requests containing sensitive information. "
|
||||
"Please rephrase your question without including passwords, secrets, or other "
|
||||
"sensitive data.",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Set terminate flag to stop execution
|
||||
context.terminate = True
|
||||
return
|
||||
|
||||
# Continue to next middleware or AI execution
|
||||
await next(context)
|
||||
|
||||
|
||||
async def class_based_chat_middleware() -> None:
|
||||
"""Demonstrate class-based middleware at agent level."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Class-based Chat Middleware (Agent Level)")
|
||||
print("=" * 60)
|
||||
|
||||
# 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="EnhancedChatAgent",
|
||||
instructions="You are a helpful AI assistant.",
|
||||
# Register class-based middleware at agent level (applies to all runs)
|
||||
middleware=[InputObserverMiddleware()],
|
||||
tools=get_weather,
|
||||
) as agent,
|
||||
):
|
||||
query = "What's the weather in Seattle?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Final Response: {result.text if result.text else 'No response'}")
|
||||
|
||||
|
||||
async def function_based_chat_middleware() -> None:
|
||||
"""Demonstrate function-based middleware at agent level."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Function-based Chat Middleware (Agent Level)")
|
||||
print("=" * 60)
|
||||
|
||||
async with (
|
||||
AzureCliCredential() as credential,
|
||||
AzureAIAgentClient(credential=credential).as_agent(
|
||||
name="FunctionMiddlewareAgent",
|
||||
instructions="You are a helpful AI assistant.",
|
||||
# Register function-based middleware at agent level
|
||||
middleware=[security_and_override_middleware],
|
||||
) as agent,
|
||||
):
|
||||
# Scenario with normal query
|
||||
print("\n--- Scenario 1: Normal Query ---")
|
||||
query = "Hello, how are you?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Final Response: {result.text if result.text else 'No response'}")
|
||||
|
||||
# Scenario with security violation
|
||||
print("\n--- Scenario 2: Security Violation ---")
|
||||
query = "What is my password for this account?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Final Response: {result.text if result.text else 'No response'}")
|
||||
|
||||
|
||||
async def run_level_middleware() -> None:
|
||||
"""Demonstrate middleware registration at run level."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Run-level Chat Middleware")
|
||||
print("=" * 60)
|
||||
|
||||
async with (
|
||||
AzureCliCredential() as credential,
|
||||
AzureAIAgentClient(credential=credential).as_agent(
|
||||
name="RunLevelAgent",
|
||||
instructions="You are a helpful AI assistant.",
|
||||
tools=get_weather,
|
||||
# No middleware at agent level
|
||||
) as agent,
|
||||
):
|
||||
# Scenario 1: Run without any middleware
|
||||
print("\n--- Scenario 1: No Middleware ---")
|
||||
query = "What's the weather in Tokyo?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Response: {result.text if result.text else 'No response'}")
|
||||
|
||||
# Scenario 2: Run with specific middleware for this call only (both enhancement and security)
|
||||
print("\n--- Scenario 2: With Run-level Middleware ---")
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(
|
||||
query,
|
||||
middleware=[
|
||||
InputObserverMiddleware(replacement="What's the weather in Madrid?"),
|
||||
security_and_override_middleware,
|
||||
],
|
||||
)
|
||||
print(f"Response: {result.text if result.text else 'No response'}")
|
||||
|
||||
# Scenario 3: Security test with run-level middleware
|
||||
print("\n--- Scenario 3: Security Test with Run-level Middleware ---")
|
||||
query = "Can you help me with my secret API key?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(
|
||||
query,
|
||||
middleware=[security_and_override_middleware],
|
||||
)
|
||||
print(f"Response: {result.text if result.text else 'No response'}")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Run all chat middleware examples."""
|
||||
print("Chat Middleware Examples")
|
||||
print("========================")
|
||||
|
||||
await class_based_chat_middleware()
|
||||
await function_based_chat_middleware()
|
||||
await run_level_middleware()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,125 @@
|
||||
# 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 (
|
||||
AgentMiddleware,
|
||||
AgentResponse,
|
||||
AgentRunContext,
|
||||
ChatMessage,
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddleware,
|
||||
Role,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Class-based Middleware 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.
|
||||
"""
|
||||
|
||||
|
||||
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: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], 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=[
|
||||
ChatMessage(role=Role.ASSISTANT, text="Detected sensitive information, the request is blocked.")
|
||||
]
|
||||
)
|
||||
# Simply don't call next() to prevent execution
|
||||
return
|
||||
|
||||
print("[SecurityAgentMiddleware] Security check passed.")
|
||||
await next(context)
|
||||
|
||||
|
||||
class LoggingFunctionMiddleware(FunctionMiddleware):
|
||||
"""Function middleware that logs function calls."""
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
function_name = context.function.name
|
||||
print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
await next(context)
|
||||
|
||||
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 Middleware 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())
|
||||
@@ -0,0 +1,87 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
|
||||
from agent_framework import (
|
||||
agent_middleware,
|
||||
function_middleware,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
|
||||
"""
|
||||
Decorator Middleware Example
|
||||
|
||||
This sample demonstrates how to use @agent_middleware and @function_middleware decorators
|
||||
to explicitly mark middleware functions without requiring type annotations.
|
||||
|
||||
The framework supports the following middleware detection scenarios:
|
||||
|
||||
1. Both decorator and parameter type specified:
|
||||
- Validates that they match (e.g., @agent_middleware with AgentRunContext)
|
||||
- Throws exception if they don't match for safety
|
||||
|
||||
2. Only decorator specified:
|
||||
- Relies on decorator to determine middleware type
|
||||
- No type annotations needed - framework handles context types automatically
|
||||
|
||||
3. Only parameter type specified:
|
||||
- Uses type annotations (AgentRunContext, FunctionInvocationContext) for detection
|
||||
|
||||
4. Neither decorator nor parameter type specified:
|
||||
- Throws exception requiring either decorator or type annotation
|
||||
- Prevents ambiguous middleware that can't be properly classified
|
||||
|
||||
Key benefits of decorator approach:
|
||||
- No type annotations needed (simpler syntax)
|
||||
- Explicit middleware type declaration
|
||||
- Clear intent in code
|
||||
- Prevents type mismatches
|
||||
"""
|
||||
|
||||
|
||||
def get_current_time() -> str:
|
||||
"""Get the current time."""
|
||||
return f"Current time is {datetime.datetime.now().strftime('%H:%M:%S')}"
|
||||
|
||||
|
||||
@agent_middleware # Decorator marks this as agent middleware - no type annotations needed
|
||||
async def simple_agent_middleware(context, next): # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
|
||||
"""Agent middleware that runs before and after agent execution."""
|
||||
print("[Agent Middleware] Before agent execution")
|
||||
await next(context)
|
||||
print("[Agent Middleware] After agent execution")
|
||||
|
||||
|
||||
@function_middleware # Decorator marks this as function middleware - no type annotations needed
|
||||
async def simple_function_middleware(context, next): # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
|
||||
"""Function middleware that runs before and after function calls."""
|
||||
print(f"[Function Middleware] Before calling: {context.function.name}") # type: ignore
|
||||
await next(context)
|
||||
print(f"[Function Middleware] After calling: {context.function.name}") # type: ignore
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating decorator-based middleware."""
|
||||
print("=== Decorator Middleware 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="TimeAgent",
|
||||
instructions="You are a helpful time assistant. Call get_current_time when asked about time.",
|
||||
tools=get_current_time,
|
||||
middleware=[simple_agent_middleware, simple_function_middleware],
|
||||
) as agent,
|
||||
):
|
||||
query = "What time is it?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import FunctionInvocationContext
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Exception Handling with Middleware
|
||||
|
||||
This sample demonstrates how to use middleware for centralized exception handling in function calls.
|
||||
The example shows:
|
||||
|
||||
- How to catch exceptions thrown by functions and provide graceful error responses
|
||||
- Overriding function results when errors occur to provide user-friendly messages
|
||||
- Using middleware to implement retry logic, fallback mechanisms, or error reporting
|
||||
|
||||
The middleware catches TimeoutError from an unstable data service and replaces it with
|
||||
a helpful message for the user, preventing raw exceptions from reaching the end user.
|
||||
"""
|
||||
|
||||
|
||||
def unstable_data_service(
|
||||
query: Annotated[str, Field(description="The data query to execute.")],
|
||||
) -> str:
|
||||
"""A simulated data service that sometimes throws exceptions."""
|
||||
# Simulate failure
|
||||
raise TimeoutError("Data service request timed out")
|
||||
|
||||
|
||||
async def exception_handling_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
function_name = context.function.name
|
||||
|
||||
try:
|
||||
print(f"[ExceptionHandlingMiddleware] Executing function: {function_name}")
|
||||
await next(context)
|
||||
print(f"[ExceptionHandlingMiddleware] Function {function_name} completed successfully.")
|
||||
except TimeoutError as e:
|
||||
print(f"[ExceptionHandlingMiddleware] Caught TimeoutError: {e}")
|
||||
# Override function result to provide custom message in response.
|
||||
context.result = (
|
||||
"Request Timeout: The data service is taking longer than expected to respond.",
|
||||
"Respond with message - 'Sorry for the inconvenience, please try again later.'",
|
||||
)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating exception handling with middleware."""
|
||||
print("=== Exception Handling Middleware 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="DataAgent",
|
||||
instructions="You are a helpful data assistant. Use the data service tool to fetch information for users.",
|
||||
tools=unstable_data_service,
|
||||
middleware=[exception_handling_middleware],
|
||||
) as agent,
|
||||
):
|
||||
query = "Get user statistics"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,109 @@
|
||||
# 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 (
|
||||
AgentRunContext,
|
||||
FunctionInvocationContext,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Function-based Middleware Example
|
||||
|
||||
This sample demonstrates how to implement middleware using simple async functions instead of classes.
|
||||
The example includes:
|
||||
|
||||
- Security middleware that validates agent requests for sensitive information
|
||||
- Logging middleware that tracks function execution timing and parameters
|
||||
- Performance monitoring to measure execution duration
|
||||
|
||||
Function-based middleware is ideal for simple, stateless operations and provides a more
|
||||
lightweight approach compared to class-based middleware. Both agent and function middleware
|
||||
can be implemented as async functions that accept context and next parameters.
|
||||
"""
|
||||
|
||||
|
||||
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."
|
||||
|
||||
|
||||
async def security_agent_middleware(
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Agent middleware that checks for security violations."""
|
||||
# Check for potential security violations in the query
|
||||
# For this example, we'll check 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.")
|
||||
# Simply don't call next() to prevent execution
|
||||
return
|
||||
|
||||
print("[SecurityAgentMiddleware] Security check passed.")
|
||||
await next(context)
|
||||
|
||||
|
||||
async def logging_function_middleware(
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Function middleware that logs function calls."""
|
||||
function_name = context.function.name
|
||||
print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
await next(context)
|
||||
|
||||
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 function-based middleware."""
|
||||
print("=== Function-based Middleware 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=[security_agent_middleware, logging_function_middleware],
|
||||
) as agent,
|
||||
):
|
||||
# Test with normal query
|
||||
print("\n--- Normal Query ---")
|
||||
query = "What's the weather like in Tokyo?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}\n")
|
||||
|
||||
# Test with security violation
|
||||
print("--- Security Test ---")
|
||||
query = "What's the secret weather password?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,177 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from random import randint
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import (
|
||||
AgentMiddleware,
|
||||
AgentResponse,
|
||||
AgentRunContext,
|
||||
ChatMessage,
|
||||
Role,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Middleware Termination Example
|
||||
|
||||
This sample demonstrates how middleware can terminate execution using the `context.terminate` flag.
|
||||
The example includes:
|
||||
|
||||
- PreTerminationMiddleware: Terminates execution before calling next() to prevent agent processing
|
||||
- PostTerminationMiddleware: Allows processing to complete but terminates further execution
|
||||
|
||||
This is useful for implementing security checks, rate limiting, or early exit conditions.
|
||||
"""
|
||||
|
||||
|
||||
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 PreTerminationMiddleware(AgentMiddleware):
|
||||
"""Middleware that terminates execution before calling the agent."""
|
||||
|
||||
def __init__(self, blocked_words: list[str]):
|
||||
self.blocked_words = [word.lower() for word in blocked_words]
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Check if the user message contains any blocked words
|
||||
last_message = context.messages[-1] if context.messages else None
|
||||
if last_message and last_message.text:
|
||||
query = last_message.text.lower()
|
||||
for blocked_word in self.blocked_words:
|
||||
if blocked_word in query:
|
||||
print(f"[PreTerminationMiddleware] Blocked word '{blocked_word}' detected. Terminating request.")
|
||||
|
||||
# Set a custom response
|
||||
context.result = AgentResponse(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
text=(
|
||||
f"Sorry, I cannot process requests containing '{blocked_word}'. "
|
||||
"Please rephrase your question."
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# Set terminate flag to prevent further processing
|
||||
context.terminate = True
|
||||
break
|
||||
|
||||
await next(context)
|
||||
|
||||
|
||||
class PostTerminationMiddleware(AgentMiddleware):
|
||||
"""Middleware that allows processing but terminates after reaching max responses across multiple runs."""
|
||||
|
||||
def __init__(self, max_responses: int = 1):
|
||||
self.max_responses = max_responses
|
||||
self.response_count = 0
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
print(f"[PostTerminationMiddleware] Processing request (response count: {self.response_count})")
|
||||
|
||||
# Check if we should terminate before processing
|
||||
if self.response_count >= self.max_responses:
|
||||
print(
|
||||
f"[PostTerminationMiddleware] Maximum responses ({self.max_responses}) reached. "
|
||||
"Terminating further processing."
|
||||
)
|
||||
context.terminate = True
|
||||
|
||||
# Allow the agent to process normally
|
||||
await next(context)
|
||||
|
||||
# Increment response count after processing
|
||||
self.response_count += 1
|
||||
|
||||
|
||||
async def pre_termination_middleware() -> None:
|
||||
"""Demonstrate pre-termination middleware that blocks requests with certain words."""
|
||||
print("\n--- Example 1: Pre-termination Middleware ---")
|
||||
async with (
|
||||
AzureCliCredential() as credential,
|
||||
AzureAIAgentClient(credential=credential).as_agent(
|
||||
name="WeatherAgent",
|
||||
instructions="You are a helpful weather assistant.",
|
||||
tools=get_weather,
|
||||
middleware=[PreTerminationMiddleware(blocked_words=["bad", "inappropriate"])],
|
||||
) as agent,
|
||||
):
|
||||
# Test with normal query
|
||||
print("\n1. Normal query:")
|
||||
query = "What's the weather like in Seattle?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text}")
|
||||
|
||||
# Test with blocked word
|
||||
print("\n2. Query with blocked word:")
|
||||
query = "What's the bad weather in New York?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text}")
|
||||
|
||||
|
||||
async def post_termination_middleware() -> None:
|
||||
"""Demonstrate post-termination middleware that limits responses across multiple runs."""
|
||||
print("\n--- Example 2: Post-termination Middleware ---")
|
||||
async with (
|
||||
AzureCliCredential() as credential,
|
||||
AzureAIAgentClient(credential=credential).as_agent(
|
||||
name="WeatherAgent",
|
||||
instructions="You are a helpful weather assistant.",
|
||||
tools=get_weather,
|
||||
middleware=[PostTerminationMiddleware(max_responses=1)],
|
||||
) as agent,
|
||||
):
|
||||
# First run (should work)
|
||||
print("\n1. First run:")
|
||||
query = "What's the weather in Paris?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text}")
|
||||
|
||||
# Second run (should be terminated by middleware)
|
||||
print("\n2. Second run (should be terminated):")
|
||||
query = "What about the weather in London?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response (terminated)'}")
|
||||
|
||||
# Third run (should also be terminated)
|
||||
print("\n3. Third run (should also be terminated):")
|
||||
query = "And New York?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response (terminated)'}")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating middleware termination functionality."""
|
||||
print("=== Middleware Termination Example ===")
|
||||
await pre_termination_middleware()
|
||||
await post_termination_middleware()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable
|
||||
from random import randint
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import (
|
||||
AgentResponse,
|
||||
AgentResponseUpdate,
|
||||
AgentRunContext,
|
||||
ChatMessage,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Result Override with Middleware (Regular and Streaming)
|
||||
|
||||
This sample demonstrates how to use middleware to intercept and modify function results
|
||||
after execution, supporting both regular and streaming agent responses. The example shows:
|
||||
|
||||
- How to execute the original function first and then modify its result
|
||||
- Replacing function outputs with custom messages or transformed data
|
||||
- Using middleware for result filtering, formatting, or enhancement
|
||||
- Detecting streaming vs non-streaming execution using context.is_streaming
|
||||
- Overriding streaming results with custom async generators
|
||||
|
||||
The weather override middleware lets the original weather function execute normally,
|
||||
then replaces its result with a custom "perfect weather" message. For streaming responses,
|
||||
it creates a custom async generator that yields the override message in chunks.
|
||||
"""
|
||||
|
||||
|
||||
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."
|
||||
|
||||
|
||||
async def weather_override_middleware(
|
||||
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
"""Middleware that overrides weather results for both streaming and non-streaming cases."""
|
||||
|
||||
# Let the original agent execution complete first
|
||||
await next(context)
|
||||
|
||||
# Check if there's a result to override (agent called weather function)
|
||||
if context.result is not None:
|
||||
# Create custom weather message
|
||||
chunks = [
|
||||
"Weather Advisory - ",
|
||||
"due to special atmospheric conditions, ",
|
||||
"all locations are experiencing perfect weather today! ",
|
||||
"Temperature is a comfortable 22°C with gentle breezes. ",
|
||||
"Perfect day for outdoor activities!",
|
||||
]
|
||||
|
||||
if context.is_streaming:
|
||||
# For streaming: create an async generator that yields chunks
|
||||
async def override_stream() -> AsyncIterable[AgentResponseUpdate]:
|
||||
for chunk in chunks:
|
||||
yield AgentResponseUpdate(contents=[TextContent(text=chunk)])
|
||||
|
||||
context.result = override_stream()
|
||||
else:
|
||||
# For non-streaming: just replace with the string message
|
||||
custom_message = "".join(chunks)
|
||||
context.result = AgentResponse(messages=[ChatMessage(role=Role.ASSISTANT, text=custom_message)])
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating result override with middleware for both streaming and non-streaming."""
|
||||
print("=== Result Override Middleware 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. Use the weather tool to get current conditions.",
|
||||
tools=get_weather,
|
||||
middleware=[weather_override_middleware],
|
||||
) as agent,
|
||||
):
|
||||
# Non-streaming example
|
||||
print("\n--- Non-streaming Example ---")
|
||||
query = "What's the weather like in Seattle?"
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result}")
|
||||
|
||||
# Streaming example
|
||||
print("\n--- Streaming Example ---")
|
||||
query = "What's the weather like in Portland?"
|
||||
print(f"User: {query}")
|
||||
print("Agent: ", end="", flush=True)
|
||||
async for chunk in agent.run_stream(query):
|
||||
if chunk.text:
|
||||
print(chunk.text, end="", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,456 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import FunctionInvocationContext, ai_function, function_middleware
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Runtime Context Delegation Patterns
|
||||
|
||||
This sample demonstrates different patterns for passing runtime context (API tokens,
|
||||
session data, etc.) to tools and sub-agents.
|
||||
|
||||
Patterns Demonstrated:
|
||||
|
||||
1. **Pattern 1: Single Agent with Middleware & Closure** (Lines 130-180)
|
||||
- Best for: Single agent with multiple tools
|
||||
- How: Middleware stores kwargs in container, tools access via closure
|
||||
- Pros: Simple, explicit state management
|
||||
- Cons: Requires container instance per agent
|
||||
|
||||
2. **Pattern 2: Hierarchical Agents with kwargs Propagation** (Lines 190-240)
|
||||
- Best for: Parent-child agent delegation with as_tool()
|
||||
- How: kwargs automatically propagate through as_tool() wrapper
|
||||
- Pros: Automatic, works with nested delegation, clean separation
|
||||
- Cons: None - this is the recommended pattern for hierarchical agents
|
||||
|
||||
3. **Pattern 3: Mixed - Hierarchical with Middleware** (Lines 250-300)
|
||||
- Best for: Complex scenarios needing both delegation and state management
|
||||
- How: Combines automatic kwargs propagation with middleware processing
|
||||
- Pros: Maximum flexibility, can transform/validate context at each level
|
||||
- Cons: More complex setup
|
||||
|
||||
Key Concepts:
|
||||
- Runtime Context: Session-specific data like API tokens, user IDs, tenant info
|
||||
- Middleware: Intercepts function calls to access/modify kwargs
|
||||
- Closure: Functions capturing variables from outer scope
|
||||
- kwargs Propagation: Automatic forwarding of runtime context through delegation chains
|
||||
"""
|
||||
|
||||
|
||||
class SessionContextContainer:
|
||||
"""Container for runtime session context accessible via closure."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Initialize with None values for runtime context."""
|
||||
self.api_token: str | None = None
|
||||
self.user_id: str | None = None
|
||||
self.session_metadata: dict[str, str] = {}
|
||||
|
||||
async def inject_context_middleware(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Middleware that extracts runtime context from kwargs and stores in container.
|
||||
|
||||
This middleware runs before tool execution and makes runtime context
|
||||
available to tools via the container instance.
|
||||
"""
|
||||
# Extract runtime context from kwargs
|
||||
self.api_token = context.kwargs.get("api_token")
|
||||
self.user_id = context.kwargs.get("user_id")
|
||||
self.session_metadata = context.kwargs.get("session_metadata", {})
|
||||
|
||||
# Log what we captured (for demonstration)
|
||||
if self.api_token or self.user_id:
|
||||
print("[Middleware] Captured runtime context:")
|
||||
print(f" - API Token: {'[PRESENT]' if self.api_token else '[NOT PROVIDED]'}")
|
||||
print(f" - User ID: {'[PRESENT]' if self.user_id else '[NOT PROVIDED]'}")
|
||||
print(f" - Session Metadata Keys: {list(self.session_metadata.keys())}")
|
||||
|
||||
# Continue to tool execution
|
||||
await next(context)
|
||||
|
||||
|
||||
# Create a container instance that will be shared via closure
|
||||
runtime_context = SessionContextContainer()
|
||||
|
||||
|
||||
@ai_function
|
||||
async def send_email(
|
||||
to: Annotated[str, Field(description="Recipient email address")],
|
||||
subject: Annotated[str, Field(description="Email subject line")],
|
||||
body: Annotated[str, Field(description="Email body content")],
|
||||
) -> str:
|
||||
"""Send an email using authenticated API (simulated).
|
||||
|
||||
This function accesses runtime context (API token, user ID) via closure
|
||||
from the runtime_context container.
|
||||
"""
|
||||
# Access runtime context via closure
|
||||
token = runtime_context.api_token
|
||||
user_id = runtime_context.user_id
|
||||
tenant = runtime_context.session_metadata.get("tenant", "unknown")
|
||||
|
||||
print("\n[send_email] Executing with runtime context:")
|
||||
print(f" - Token: {'[PRESENT]' if token else '[NOT PROVIDED]'}")
|
||||
print(f" - User ID: {'[PRESENT]' if user_id else '[NOT PROVIDED]'}")
|
||||
print(f" - Tenant: {'[PRESENT]' if tenant and tenant != 'unknown' else '[NOT PROVIDED]'}")
|
||||
print(" - Recipient count: 1")
|
||||
print(f" - Subject length: {len(subject)} chars")
|
||||
|
||||
# Simulate API call with authentication
|
||||
if not token:
|
||||
return "ERROR: No API token provided - cannot send email"
|
||||
|
||||
# Simulate sending email
|
||||
return f"Email sent to {to} from user {user_id} (tenant: {tenant}). Subject: '{subject}'"
|
||||
|
||||
|
||||
@ai_function
|
||||
async def send_notification(
|
||||
message: Annotated[str, Field(description="Notification message to send")],
|
||||
priority: Annotated[str, Field(description="Priority level: low, medium, high")] = "medium",
|
||||
) -> str:
|
||||
"""Send a push notification using authenticated API (simulated).
|
||||
|
||||
This function accesses runtime context via closure from runtime_context.
|
||||
"""
|
||||
token = runtime_context.api_token
|
||||
user_id = runtime_context.user_id
|
||||
|
||||
print("\n[send_notification] Executing with runtime context:")
|
||||
print(f" - Token: {'[PRESENT]' if token else '[NOT PROVIDED]'}")
|
||||
print(f" - User ID: {'[PRESENT]' if user_id else '[NOT PROVIDED]'}")
|
||||
print(f" - Message length: {len(message)} chars")
|
||||
print(f" - Priority: {priority}")
|
||||
|
||||
if not token:
|
||||
return "ERROR: No API token provided - cannot send notification"
|
||||
|
||||
return f"Notification sent to user {user_id} with priority {priority}: {message}"
|
||||
|
||||
|
||||
async def pattern_1_single_agent_with_closure() -> None:
|
||||
"""Pattern 1: Single agent with middleware and closure for runtime context."""
|
||||
print("\n" + "=" * 70)
|
||||
print("PATTERN 1: Single Agent with Middleware & Closure")
|
||||
print("=" * 70)
|
||||
print("Use case: Single agent with multiple tools sharing runtime context")
|
||||
print()
|
||||
|
||||
client = OpenAIChatClient(model_id="gpt-4o-mini")
|
||||
|
||||
# Create agent with both tools and shared context via middleware
|
||||
communication_agent = client.as_agent(
|
||||
name="communication_agent",
|
||||
instructions=(
|
||||
"You are a communication assistant that can send emails and notifications. "
|
||||
"Use send_email for email tasks and send_notification for notification tasks."
|
||||
),
|
||||
tools=[send_email, send_notification],
|
||||
# Both tools share the same context container via middleware
|
||||
middleware=[runtime_context.inject_context_middleware],
|
||||
)
|
||||
|
||||
# Test 1: Send email with runtime context
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 1: Email with Runtime Context")
|
||||
print("=" * 70)
|
||||
|
||||
user_query = (
|
||||
"Send an email to john@example.com with subject 'Meeting Tomorrow' and body 'Don't forget our 2pm meeting.'"
|
||||
)
|
||||
print(f"\nUser: {user_query}")
|
||||
|
||||
result1 = await communication_agent.run(
|
||||
user_query,
|
||||
# Runtime context passed as kwargs
|
||||
api_token="sk-test-token-xyz-789",
|
||||
user_id="user-12345",
|
||||
session_metadata={"tenant": "acme-corp", "region": "us-west"},
|
||||
)
|
||||
|
||||
print(f"\nAgent: {result1.text}")
|
||||
|
||||
# Test 2: Send notification with different runtime context
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 2: Notification with Different Runtime Context")
|
||||
print("=" * 70)
|
||||
|
||||
user_query2 = "Send a high priority notification saying 'Your order has shipped!'"
|
||||
print(f"\nUser: {user_query2}")
|
||||
|
||||
result2 = await communication_agent.run(
|
||||
user_query2,
|
||||
# Different runtime context for this request
|
||||
api_token="sk-prod-token-abc-456",
|
||||
user_id="user-67890",
|
||||
session_metadata={"tenant": "store-inc", "region": "eu-central"},
|
||||
)
|
||||
|
||||
print(f"\nAgent: {result2.text}")
|
||||
|
||||
# Test 3: Both email and notification in one request
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 3: Multiple Tools in One Request")
|
||||
print("=" * 70)
|
||||
|
||||
user_query3 = (
|
||||
"Send an email to alice@example.com about the new feature launch "
|
||||
"and also send a notification to remind about the team meeting."
|
||||
)
|
||||
print(f"\nUser: {user_query3}")
|
||||
|
||||
result3 = await communication_agent.run(
|
||||
user_query3,
|
||||
api_token="sk-dev-token-def-123",
|
||||
user_id="user-11111",
|
||||
session_metadata={"tenant": "dev-team", "region": "us-east"},
|
||||
)
|
||||
|
||||
print(f"\nAgent: {result3.text}")
|
||||
|
||||
# Test 4: Missing context - show error handling
|
||||
print("\n" + "=" * 70)
|
||||
print("TEST 4: Missing Runtime Context (Error Case)")
|
||||
print("=" * 70)
|
||||
|
||||
user_query4 = "Send an email to test@example.com with subject 'Test'"
|
||||
print(f"\nUser: {user_query4}")
|
||||
print("Note: Running WITHOUT api_token to demonstrate error handling")
|
||||
|
||||
result4 = await communication_agent.run(
|
||||
user_query4,
|
||||
# Missing api_token - tools should handle gracefully
|
||||
user_id="user-22222",
|
||||
)
|
||||
|
||||
print(f"\nAgent: {result4.text}")
|
||||
|
||||
print("\n✓ Pattern 1 complete - Middleware & closure pattern works for single agents")
|
||||
|
||||
|
||||
# Pattern 2: Hierarchical agents with automatic kwargs propagation
|
||||
# ================================================================
|
||||
|
||||
|
||||
# Create tools for sub-agents (these will use kwargs propagation)
|
||||
@ai_function
|
||||
async def send_email_v2(
|
||||
to: Annotated[str, Field(description="Recipient email")],
|
||||
subject: Annotated[str, Field(description="Subject")],
|
||||
body: Annotated[str, Field(description="Body")],
|
||||
) -> str:
|
||||
"""Send email - demonstrates kwargs propagation pattern."""
|
||||
# In this pattern, we can create a middleware to access kwargs
|
||||
# But for simplicity, we'll just simulate the operation
|
||||
return f"Email sent to {to} with subject '{subject}'"
|
||||
|
||||
|
||||
@ai_function
|
||||
async def send_sms(
|
||||
phone: Annotated[str, Field(description="Phone number")],
|
||||
message: Annotated[str, Field(description="SMS message")],
|
||||
) -> str:
|
||||
"""Send SMS message."""
|
||||
return f"SMS sent to {phone}: {message}"
|
||||
|
||||
|
||||
async def pattern_2_hierarchical_with_kwargs_propagation() -> None:
|
||||
"""Pattern 2: Hierarchical agents with automatic kwargs propagation through as_tool()."""
|
||||
print("\n" + "=" * 70)
|
||||
print("PATTERN 2: Hierarchical Agents with kwargs Propagation")
|
||||
print("=" * 70)
|
||||
print("Use case: Parent agent delegates to specialized sub-agents")
|
||||
print("Feature: Runtime kwargs automatically propagate through as_tool()")
|
||||
print()
|
||||
|
||||
# Track kwargs at each level
|
||||
email_agent_kwargs: dict[str, object] = {}
|
||||
sms_agent_kwargs: dict[str, object] = {}
|
||||
|
||||
@function_middleware
|
||||
async def email_kwargs_tracker(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
email_agent_kwargs.update(context.kwargs)
|
||||
print(f"[EmailAgent] Received runtime context: {list(context.kwargs.keys())}")
|
||||
await next(context)
|
||||
|
||||
@function_middleware
|
||||
async def sms_kwargs_tracker(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
sms_agent_kwargs.update(context.kwargs)
|
||||
print(f"[SMSAgent] Received runtime context: {list(context.kwargs.keys())}")
|
||||
await next(context)
|
||||
|
||||
client = OpenAIChatClient(model_id="gpt-4o-mini")
|
||||
|
||||
# Create specialized sub-agents
|
||||
email_agent = client.as_agent(
|
||||
name="email_agent",
|
||||
instructions="You send emails using the send_email_v2 tool.",
|
||||
tools=[send_email_v2],
|
||||
middleware=[email_kwargs_tracker],
|
||||
)
|
||||
|
||||
sms_agent = client.as_agent(
|
||||
name="sms_agent",
|
||||
instructions="You send SMS messages using the send_sms tool.",
|
||||
tools=[send_sms],
|
||||
middleware=[sms_kwargs_tracker],
|
||||
)
|
||||
|
||||
# Create coordinator that delegates to sub-agents
|
||||
coordinator = client.as_agent(
|
||||
name="coordinator",
|
||||
instructions=(
|
||||
"You coordinate communication tasks. "
|
||||
"Use email_sender for emails and sms_sender for SMS. "
|
||||
"Delegate to the appropriate specialized agent."
|
||||
),
|
||||
tools=[
|
||||
email_agent.as_tool(
|
||||
name="email_sender",
|
||||
description="Send emails to recipients",
|
||||
arg_name="task",
|
||||
),
|
||||
sms_agent.as_tool(
|
||||
name="sms_sender",
|
||||
description="Send SMS messages",
|
||||
arg_name="task",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
# Test: Runtime context propagates automatically
|
||||
print("Test: Send email with runtime context\n")
|
||||
await coordinator.run(
|
||||
"Send an email to john@example.com with subject 'Meeting' and body 'See you at 2pm'",
|
||||
api_token="secret-token-abc",
|
||||
user_id="user-999",
|
||||
tenant_id="tenant-acme",
|
||||
)
|
||||
|
||||
print(f"\n[Verification] EmailAgent received kwargs keys: {list(email_agent_kwargs.keys())}")
|
||||
print(f" - api_token: {'[PRESENT]' if email_agent_kwargs.get('api_token') else '[NOT PROVIDED]'}")
|
||||
print(f" - user_id: {'[PRESENT]' if email_agent_kwargs.get('user_id') else '[NOT PROVIDED]'}")
|
||||
print(f" - tenant_id: {'[PRESENT]' if email_agent_kwargs.get('tenant_id') else '[NOT PROVIDED]'}")
|
||||
|
||||
print("\n✓ Pattern 2 complete - kwargs automatically propagate through as_tool()")
|
||||
|
||||
|
||||
# Pattern 3: Mixed pattern - hierarchical with middleware processing
|
||||
# ===================================================================
|
||||
|
||||
|
||||
class AuthContextMiddleware:
|
||||
"""Middleware that validates and transforms runtime context."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.validated_tokens: list[str] = []
|
||||
|
||||
async def validate_and_track(
|
||||
self, context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
"""Validate API token and track usage."""
|
||||
api_token = context.kwargs.get("api_token")
|
||||
|
||||
if api_token:
|
||||
# Simulate token validation
|
||||
if api_token.startswith("valid-"):
|
||||
print("[AuthMiddleware] Token validated successfully")
|
||||
self.validated_tokens.append(api_token)
|
||||
else:
|
||||
print("[AuthMiddleware] Token validation failed")
|
||||
# Could set context.terminate = True to block execution
|
||||
else:
|
||||
print("[AuthMiddleware] No API token provided")
|
||||
|
||||
await next(context)
|
||||
|
||||
|
||||
@ai_function
|
||||
async def protected_operation(operation: Annotated[str, Field(description="Operation to perform")]) -> str:
|
||||
"""Protected operation that requires authentication."""
|
||||
return f"Executed protected operation: {operation}"
|
||||
|
||||
|
||||
async def pattern_3_hierarchical_with_middleware() -> None:
|
||||
"""Pattern 3: Hierarchical agents with middleware processing at each level."""
|
||||
print("\n" + "=" * 70)
|
||||
print("PATTERN 3: Hierarchical with Middleware Processing")
|
||||
print("=" * 70)
|
||||
print("Use case: Multi-level validation/transformation of runtime context")
|
||||
print()
|
||||
|
||||
auth_middleware = AuthContextMiddleware()
|
||||
|
||||
client = OpenAIChatClient(model_id="gpt-4o-mini")
|
||||
|
||||
# Sub-agent with validation middleware
|
||||
protected_agent = client.as_agent(
|
||||
name="protected_agent",
|
||||
instructions="You perform protected operations that require authentication.",
|
||||
tools=[protected_operation],
|
||||
middleware=[auth_middleware.validate_and_track],
|
||||
)
|
||||
|
||||
# Coordinator delegates to protected agent
|
||||
coordinator = client.as_agent(
|
||||
name="coordinator",
|
||||
instructions="You coordinate protected operations. Delegate to protected_executor.",
|
||||
tools=[
|
||||
protected_agent.as_tool(
|
||||
name="protected_executor",
|
||||
description="Execute protected operations",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
# Test with valid token
|
||||
print("Test 1: Valid token\n")
|
||||
await coordinator.run(
|
||||
"Execute operation: backup_database",
|
||||
api_token="valid-token-xyz-789",
|
||||
user_id="admin-123",
|
||||
)
|
||||
|
||||
# Test with invalid token
|
||||
print("\nTest 2: Invalid token\n")
|
||||
await coordinator.run(
|
||||
"Execute operation: delete_records",
|
||||
api_token="invalid-token-bad",
|
||||
user_id="user-456",
|
||||
)
|
||||
|
||||
print(f"\n[Validation Summary] Validated tokens: {len(auth_middleware.validated_tokens)}")
|
||||
print("✓ Pattern 3 complete - Middleware can validate/transform context at each level")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Demonstrate all runtime context delegation patterns."""
|
||||
print("=" * 70)
|
||||
print("Runtime Context Delegation Patterns Demo")
|
||||
print("=" * 70)
|
||||
print()
|
||||
|
||||
# Run Pattern 1
|
||||
await pattern_1_single_agent_with_closure()
|
||||
|
||||
# Run Pattern 2
|
||||
await pattern_2_hierarchical_with_kwargs_propagation()
|
||||
|
||||
# Run Pattern 3
|
||||
await pattern_3_hierarchical_with_middleware()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,128 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from random import randint
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import (
|
||||
FunctionInvocationContext,
|
||||
)
|
||||
from agent_framework.azure import AzureAIAgentClient
|
||||
from azure.identity.aio import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Shared State Function-based Middleware Example
|
||||
|
||||
This sample demonstrates how to implement function-based middleware within a class to share state.
|
||||
The example includes:
|
||||
|
||||
- A MiddlewareContainer class with two simple function middleware methods
|
||||
- First middleware: Counts function calls and stores the count in shared state
|
||||
- Second middleware: Uses the shared count to add call numbers to function results
|
||||
|
||||
This approach shows how middleware can work together by sharing state within the same class instance.
|
||||
"""
|
||||
|
||||
|
||||
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."
|
||||
|
||||
|
||||
def get_time(
|
||||
timezone: Annotated[str, Field(description="The timezone to get the time for.")] = "UTC",
|
||||
) -> str:
|
||||
"""Get the current time for a given timezone."""
|
||||
import datetime
|
||||
|
||||
return f"The current time in {timezone} is {datetime.datetime.now().strftime('%H:%M:%S')}"
|
||||
|
||||
|
||||
class MiddlewareContainer:
|
||||
"""Container class that holds middleware functions with shared state."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Simple shared state: count function calls
|
||||
self.call_count: int = 0
|
||||
|
||||
async def call_counter_middleware(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""First middleware: increments call count in shared state."""
|
||||
# Increment the shared call count
|
||||
self.call_count += 1
|
||||
|
||||
print(f"[CallCounter] This is function call #{self.call_count}")
|
||||
|
||||
# Call the next middleware/function
|
||||
await next(context)
|
||||
|
||||
async def result_enhancer_middleware(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Second middleware: uses shared call count to enhance function results."""
|
||||
print(f"[ResultEnhancer] Current total calls so far: {self.call_count}")
|
||||
|
||||
# Call the next middleware/function
|
||||
await next(context)
|
||||
|
||||
# After function execution, enhance the result using shared state
|
||||
if context.result:
|
||||
enhanced_result = f"[Call #{self.call_count}] {context.result}"
|
||||
context.result = enhanced_result
|
||||
print("[ResultEnhancer] Enhanced result with call number")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating shared state function-based middleware."""
|
||||
print("=== Shared State Function-based Middleware Example ===")
|
||||
|
||||
# Create middleware container with shared state
|
||||
middleware_container = MiddlewareContainer()
|
||||
|
||||
# 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="UtilityAgent",
|
||||
instructions="You are a helpful assistant that can provide weather information and current time.",
|
||||
tools=[get_weather, get_time],
|
||||
# Pass both middleware functions from the same container instance
|
||||
# Order matters: counter runs first to increment count,
|
||||
# then result enhancer uses the updated count
|
||||
middleware=[
|
||||
middleware_container.call_counter_middleware,
|
||||
middleware_container.result_enhancer_middleware,
|
||||
],
|
||||
) as agent,
|
||||
):
|
||||
# Test multiple requests to see shared state in action
|
||||
queries = [
|
||||
"What's the weather like in New York?",
|
||||
"What time is it in London?",
|
||||
"What's the weather in Tokyo?",
|
||||
]
|
||||
|
||||
for i, query in enumerate(queries, 1):
|
||||
print(f"\n--- Query {i} ---")
|
||||
print(f"User: {query}")
|
||||
result = await agent.run(query)
|
||||
print(f"Agent: {result.text if result.text else 'No response'}")
|
||||
|
||||
# Display final statistics
|
||||
print("\n=== Final Statistics ===")
|
||||
print(f"Total function calls made: {middleware_container.call_count}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,99 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Annotated
|
||||
|
||||
from agent_framework import (
|
||||
AgentRunContext,
|
||||
ChatMessageStore,
|
||||
)
|
||||
from agent_framework.azure import AzureOpenAIChatClient
|
||||
from azure.identity import AzureCliCredential
|
||||
from pydantic import Field
|
||||
|
||||
"""
|
||||
Thread Behavior Middleware Example
|
||||
|
||||
This sample demonstrates how middleware can access and track thread state across multiple agent runs.
|
||||
The example shows:
|
||||
|
||||
- How AgentRunContext.thread property behaves across multiple runs
|
||||
- How middleware can access conversation history through the thread
|
||||
- The timing of when thread messages are populated (before vs after next() call)
|
||||
- How to track thread state changes across runs
|
||||
|
||||
Key behaviors demonstrated:
|
||||
1. First run: context.messages is populated, context.thread is initially empty (before next())
|
||||
2. After next(): thread contains input message + response from agent
|
||||
3. Second run: context.messages contains only current input, thread contains previous history
|
||||
4. After next(): thread contains full conversation history (all previous + current messages)
|
||||
"""
|
||||
|
||||
|
||||
def get_weather(
|
||||
location: Annotated[str, Field(description="The location to get the weather for.")],
|
||||
) -> str:
|
||||
"""Get the weather for a given location."""
|
||||
from random import randint
|
||||
|
||||
conditions = ["sunny", "cloudy", "rainy", "stormy"]
|
||||
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
|
||||
|
||||
|
||||
async def thread_tracking_middleware(
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Middleware that tracks and logs thread behavior across runs."""
|
||||
thread_messages = []
|
||||
if context.thread and context.thread.message_store:
|
||||
thread_messages = await context.thread.message_store.list_messages()
|
||||
|
||||
print(f"[Middleware pre-execution] Current input messages: {len(context.messages)}")
|
||||
print(f"[Middleware pre-execution] Thread history messages: {len(thread_messages)}")
|
||||
|
||||
# Call next to execute the agent
|
||||
await next(context)
|
||||
|
||||
# Check thread state after agent execution
|
||||
updated_thread_messages = []
|
||||
if context.thread and context.thread.message_store:
|
||||
updated_thread_messages = await context.thread.message_store.list_messages()
|
||||
|
||||
print(f"[Middleware post-execution] Updated thread messages: {len(updated_thread_messages)}")
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Example demonstrating thread behavior in middleware across multiple runs."""
|
||||
print("=== Thread Behavior Middleware Example ===")
|
||||
|
||||
# For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
|
||||
# authentication option.
|
||||
agent = AzureOpenAIChatClient(credential=AzureCliCredential()).as_agent(
|
||||
name="WeatherAgent",
|
||||
instructions="You are a helpful weather assistant.",
|
||||
tools=get_weather,
|
||||
middleware=[thread_tracking_middleware],
|
||||
# Configure agent with message store factory to persist conversation history
|
||||
chat_message_store_factory=ChatMessageStore,
|
||||
)
|
||||
|
||||
# Create a thread that will persist messages between runs
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
print("\nFirst Run:")
|
||||
query1 = "What's the weather like in Tokyo?"
|
||||
print(f"User: {query1}")
|
||||
result1 = await agent.run(query1, thread=thread)
|
||||
print(f"Agent: {result1.text}")
|
||||
|
||||
print("\nSecond Run:")
|
||||
query2 = "How about in London?"
|
||||
print(f"User: {query2}")
|
||||
result2 = await agent.run(query2, thread=thread)
|
||||
print(f"Agent: {result2.text}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user