agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added context result override

Browse Source
This commit is contained in:
Dmytro Struk
2025-09-17 15:17:31 -07:00
Unverified
parent 2df0e289e8
commit 89eeb58bdc
5 changed files with 1300 additions and 596 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_middleware.py
+67 -16
View File
@@ -32,6 +32,10 @@ class AgentRunContext:
messages: The messages being sent to the agent.
is_streaming: Whether this is a streaming invocation.
metadata: Metadata dictionary for sharing data between agent middleware.
response: Agent execution response. Can be set before calling next() to override execution,
or observed after calling next() to see the actual execution result.
For non-streaming: should be AgentRunResponse
For streaming: should be AsyncIterable[AgentRunResponseUpdate]
"""
def __init__(
@@ -53,6 +57,7 @@ class AgentRunContext:
self.messages = messages
self.is_streaming = is_streaming
self.metadata = metadata or {}
self.response: AgentRunResponse | AsyncIterable[AgentRunResponseUpdate] | None = None
class FunctionInvocationContext:
@@ -62,6 +67,8 @@ class FunctionInvocationContext:
function: The function being invoked.
arguments: The validated arguments for the function.
metadata: Metadata dictionary for sharing data between function middleware.
result: Function execution result. Can be set before calling next() to override execution,
or observed after calling next() to see the actual execution result.
"""
def __init__(
@@ -80,6 +87,7 @@ class FunctionInvocationContext:
self.function = function
self.arguments = arguments
self.metadata = metadata or {}
self.result: Any = None
class AgentMiddleware(ABC):
@@ -96,15 +104,17 @@ class AgentMiddleware(ABC):
Args:
context: Agent invocation context containing agent, messages, and metadata.
Use context.is_streaming to determine if this is a streaming call.
Middleware can set context.should_skip=True and provide context.response
or context.response_stream to override the agent execution.
Middleware can set context.response to override execution, or observe
the actual execution result after calling next().
For non-streaming: AgentRunResponse
For streaming: AsyncIterable[AgentRunResponseUpdate]
next: Function to call the next middleware or final agent execution.
Does not return anything - all data flows through the context.
Note:
Middleware should not return anything. All data manipulation should happen
within the context object. Set context.should_skip=True and provide
context.response or context.response_stream to override execution.
within the context object. Set context.response to override execution,
or observe context.response after calling next() for actual results.
"""
...
@@ -122,15 +132,15 @@ class FunctionMiddleware(ABC):
Args:
context: Function invocation context containing function, arguments, and metadata.
Middleware can set context.should_skip=True and provide context.result
to override the function execution.
Middleware can set context.result to override execution, or observe
the actual execution result after calling next().
next: Function to call the next middleware or final function execution.
Does not return anything - all data flows through the context.
Note:
Middleware should not return anything. All data manipulation should happen
within the context object. Set context.should_skip=True and provide
context.result to override execution.
within the context object. Set context.result to override execution,
or observe context.result after calling next() for actual results.
"""
...
@@ -229,7 +239,14 @@ class AgentMiddlewarePipeline:
if index >= len(self._middlewares):
async def final_wrapper(c: AgentRunContext) -> None:
result_container["response"] = await final_handler(c)
# If response was set before calling next(), skip execution
if c.response is not None and isinstance(c.response, AgentRunResponse):
result_container["response"] = c.response
return
# Execute actual handler and populate context for observability
result = await final_handler(c)
result_container["response"] = result
c.response = result
return final_wrapper
@@ -238,14 +255,24 @@ class AgentMiddlewarePipeline:
async def current_handler(c: AgentRunContext) -> None:
await middleware.process(c, next_handler)
# After middleware execution, check if response was overridden
if c.response is not None and isinstance(c.response, AgentRunResponse):
result_container["response"] = c.response
return current_handler
first_handler = create_next_handler(0)
await first_handler(context)
# Return the response from result container
return result_container["response"]
# Return the response from result container or overridden response
if context.response is not None and isinstance(context.response, AgentRunResponse):
return context.response
# If no response was set (next() not called), return empty AgentRunResponse
response = result_container["response"]
if response is None:
return AgentRunResponse(messages=[])
return response
async def execute_stream(
self,
@@ -282,7 +309,15 @@ class AgentMiddlewarePipeline:
if index >= len(self._middlewares):
async def final_wrapper(c: AgentRunContext) -> None: # noqa: RUF029
result_container["response_stream"] = final_handler(c)
# If response was set before calling next(), skip execution
if c.response is not None and hasattr(c.response, "__aiter__"):
result_container["response_stream"] = c.response # type: ignore
return
# Execute actual handler and populate context for observability
result = final_handler(c)
result_container["response_stream"] = result
c.response = result
return final_wrapper
@@ -297,10 +332,16 @@ class AgentMiddlewarePipeline:
first_handler = create_next_handler(0)
await first_handler(context)
# Yield from the response stream in result container
# Yield from the response stream in result container or overridden response
if context.response is not None and hasattr(context.response, "__aiter__"):
async for update in context.response: # type: ignore
yield update
return
response_stream = result_container["response_stream"]
if response_stream is None:
raise RuntimeError("No response stream set after middleware execution")
# If no response stream was set (next() not called), yield nothing
return
async for update in response_stream:
yield update
@@ -366,7 +407,15 @@ class FunctionMiddlewarePipeline:
if index >= len(self._middlewares):
async def final_wrapper(c: FunctionInvocationContext) -> None:
result_container["result"] = await final_handler(c)
# If result was set before calling next(), skip execution
if c.result is not None:
result_container["result"] = c.result
return
# Execute actual handler and populate context for observability
result = await final_handler(c)
result_container["result"] = result
c.result = result
return final_wrapper
@@ -381,7 +430,9 @@ class FunctionMiddlewarePipeline:
first_handler = create_next_handler(0)
await first_handler(context)
# Return the result from result container
# Return the result from result container or overridden result
if context.result is not None:
return context.result
return result_container["result"]
@property
python/packages/main/tests/main/conftest.py
+13 -2
View File
@@ -105,6 +105,9 @@ class MockChatClient:
def __init__(self) -> None:
self.additional_properties: dict[str, Any] = {}
self.call_count: int = 0
self.responses: list[ChatResponse] = []
self.streaming_responses: list[list[ChatResponseUpdate]] = []
async def get_response(
self,
@@ -112,6 +115,9 @@ class MockChatClient:
**kwargs: Any,
) -> ChatResponse:
logger.debug(f"Running custom chat client, with: {messages=}, {kwargs=}")
self.call_count += 1
if self.responses:
return self.responses.pop(0)
return ChatResponse(messages=ChatMessage(role="assistant", text="test response"))
async def get_streaming_response(
@@ -120,8 +126,13 @@ class MockChatClient:
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
logger.debug(f"Running custom chat client stream, with: {messages=}, {kwargs=}")
yield ChatResponseUpdate(text=TextContent(text="test streaming response "), role="assistant")
yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
self.call_count += 1
if self.streaming_responses:
for update in self.streaming_responses.pop(0):
yield update
else:
yield ChatResponseUpdate(text=TextContent(text="test streaming response "), role="assistant")
yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
class MockBaseChatClient(BaseChatClient):
python/packages/main/tests/main/test_middleware.py
+213 -578
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import AsyncIterable, Awaitable, Callable, MutableSequence
from collections.abc import AsyncIterable, Awaitable, Callable
from typing import Any
from unittest.mock import MagicMock
@@ -11,16 +11,10 @@ from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentRunResponseUpdate,
ChatAgent,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
)
from agent_framework._clients import BaseChatClient
from agent_framework._middleware import (
AgentMiddleware,
AgentMiddlewarePipeline,
@@ -29,8 +23,7 @@ from agent_framework._middleware import (
FunctionMiddleware,
FunctionMiddlewarePipeline,
)
from agent_framework._tools import AIFunction, use_function_invocation
from agent_framework._types import ChatOptions
from agent_framework._tools import AIFunction
class TestAgentRunContext:
@@ -753,6 +746,217 @@ class MockFunctionArgs(BaseModel):
name: str = Field(description="Test name parameter")
class TestMiddlewareExecutionControl:
"""Test cases for middleware execution control (when next() is called vs not called)."""
async def test_agent_middleware_no_next_no_execution(self, mock_agent: AgentProtocol) -> None:
"""Test that when agent middleware doesn't call next(), no execution happens."""
class NoNextMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Don't call next() - this should prevent any execution
pass
middleware = NoNextMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
nonlocal handler_called
handler_called = True
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
# Verify no execution happened - should return empty AgentRunResponse
assert result is not None
assert isinstance(result, AgentRunResponse)
assert result.messages == [] # Empty response
assert not handler_called
assert context.response is None
async def test_agent_middleware_no_next_no_streaming_execution(self, mock_agent: AgentProtocol) -> None:
"""Test that when agent middleware doesn't call next(), no streaming execution happens."""
class NoNextStreamingMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Don't call next() - this should prevent any execution
pass
middleware = NoNextStreamingMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
nonlocal handler_called
handler_called = True
yield AgentRunResponseUpdate(contents=[TextContent(text="should not execute")])
# When middleware doesn't call next(), streaming should yield no updates
updates: list[AgentRunResponseUpdate] = []
async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
updates.append(update)
# Verify no execution happened and no updates were yielded
assert len(updates) == 0
assert not handler_called
assert context.response is None
async def test_function_middleware_no_next_no_execution(self, mock_function: AIFunction[Any, Any]) -> None:
"""Test that when function middleware doesn't call next(), no execution happens."""
class FunctionTestArgs(BaseModel):
name: str = Field(description="Test name parameter")
class NoNextFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Don't call next() - this should prevent any execution
pass
middleware = NoNextFunctionMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
arguments = FunctionTestArgs(name="test")
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
handler_called = False
async def final_handler(ctx: FunctionInvocationContext) -> str:
nonlocal handler_called
handler_called = True
return "should not execute"
result = await pipeline.execute(mock_function, arguments, context, final_handler)
# Verify no execution happened
assert result is None
assert not handler_called
assert context.result is None
async def test_multiple_middlewares_early_stop(self, mock_agent: AgentProtocol) -> None:
"""Test that when first middleware doesn't call next(), subsequent middlewares are not called."""
execution_order: list[str] = []
class FirstMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("first")
# Don't call next() - this should stop the pipeline
class SecondMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("second")
await next(context)
pipeline = AgentMiddlewarePipeline([FirstMiddleware(), SecondMiddleware()])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
nonlocal handler_called
handler_called = True
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
# Verify only first middleware was called and empty response returned
assert execution_order == ["first"]
assert result is not None
assert isinstance(result, AgentRunResponse)
assert result.messages == [] # Empty response
assert not handler_called
async def test_agent_middleware_pre_execution_override_with_next(self, mock_agent: AgentProtocol) -> None:
"""Test that middleware can override response before calling next() - this skips handler execution."""
class PreOverrideMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Set override first
context.response = AgentRunResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="pre-override response")]
)
# Then call next() to continue middleware pipeline
await next(context)
middleware = PreOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
nonlocal handler_called
handler_called = True
# This should not be called when response is pre-set
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="original response")])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
# Verify pre-override worked and handler was NOT called (because response was already set)
assert result is not None
assert result.messages[0].text == "pre-override response"
assert not handler_called
async def test_function_middleware_pre_execution_override_with_next(
self, mock_function: AIFunction[Any, Any]
) -> None:
"""Test that function middleware can override result before calling next() - this skips handler execution."""
class FunctionTestArgs(BaseModel):
name: str = Field(description="Test name parameter")
class PreOverrideFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Set override first
context.result = "pre-override result"
# Then call next() to continue middleware pipeline
await next(context)
middleware = PreOverrideFunctionMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
arguments = FunctionTestArgs(name="test")
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
handler_called = False
async def final_handler(ctx: FunctionInvocationContext) -> str:
nonlocal handler_called
handler_called = True
# This should not be called when result is pre-set
return "original result"
result = await pipeline.execute(mock_function, arguments, context, final_handler)
# Verify pre-override worked and handler was NOT called (because result was already set)
assert result == "pre-override result"
assert not handler_called
@pytest.fixture
def mock_agent() -> AgentProtocol:
"""Mock agent for testing."""
@@ -767,572 +971,3 @@ def mock_function() -> AIFunction[Any, Any]:
function = MagicMock(spec=AIFunction[Any, Any])
function.name = "test_function"
return function
@use_function_invocation
class MockChatClient(BaseChatClient):
"""Mock chat client for ChatAgent integration tests."""
call_count: int = Field(default=0)
responses: list[ChatResponse] = Field(default_factory=lambda: [])
streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=lambda: [])
async def _inner_get_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> ChatResponse:
"""Return a mock response."""
self.call_count += 1
if self.responses:
return self.responses.pop(0)
return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Mock response")])
async def _inner_get_streaming_response(
self,
*,
messages: MutableSequence[ChatMessage],
chat_options: ChatOptions,
**kwargs: Any,
) -> AsyncIterable[ChatResponseUpdate]:
"""Return mock streaming responses."""
self.call_count += 1
if self.streaming_responses:
for update in self.streaming_responses.pop(0):
yield update
else:
yield ChatResponseUpdate(contents=[TextContent(text="Mock")], role=Role.ASSISTANT)
yield ChatResponseUpdate(contents=[TextContent(text=" streaming response")], role=Role.ASSISTANT)
def service_url(self) -> str:
return "https://mock.example.com"
@pytest.fixture
def mock_chat_client() -> MockChatClient:
"""Mock chat client fixture."""
return MockChatClient()
# region ChatAgent Tests
class TestChatAgentClassBasedMiddleware:
"""Test cases for class-based middleware integration with ChatAgent."""
async def test_class_based_agent_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
"""Test class-based agent middleware with ChatAgent."""
execution_order: list[str] = []
class TrackingAgentMiddleware(AgentMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Create ChatAgent with middleware
middleware = TrackingAgentMiddleware("agent_middleware")
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "Mock response"
assert mock_chat_client.call_count == 1
# Verify middleware execution order
assert execution_order == ["agent_middleware_before", "agent_middleware_after"]
async def test_class_based_function_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
"""Test class-based function middleware with ChatAgent."""
execution_order: list[str] = []
class TrackingFunctionMiddleware(FunctionMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
middleware = TrackingFunctionMiddleware("function_middleware")
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert mock_chat_client.call_count == 1
# Note: Function middleware won't execute since no function calls are made
assert execution_order == []
class TestChatAgentFunctionBasedMiddleware:
"""Test cases for function-based middleware integration with ChatAgent."""
async def test_function_based_agent_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
"""Test function-based agent middleware with ChatAgent."""
execution_order: list[str] = []
async def tracking_agent_middleware(
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("agent_function_before")
await next(context)
execution_order.append("agent_function_after")
# Create ChatAgent with function middleware
agent = ChatAgent(chat_client=mock_chat_client, middleware=[tracking_agent_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "Mock response"
assert mock_chat_client.call_count == 1
# Verify middleware execution order
assert execution_order == ["agent_function_before", "agent_function_after"]
async def test_function_based_function_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
"""Test function-based function middleware with ChatAgent."""
execution_order: list[str] = []
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
execution_order.append("function_function_after")
# Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
agent = ChatAgent(chat_client=mock_chat_client, middleware=[tracking_function_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert mock_chat_client.call_count == 1
# Note: Function middleware won't execute since no function calls are made
assert execution_order == []
class TestChatAgentStreamingMiddleware:
"""Test cases for streaming middleware integration with ChatAgent."""
async def test_agent_middleware_with_streaming(self, mock_chat_client: MockChatClient) -> None:
"""Test agent middleware with streaming ChatAgent responses."""
execution_order: list[str] = []
streaming_flags: list[bool] = []
class StreamingTrackingMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("middleware_before")
streaming_flags.append(context.is_streaming)
await next(context)
execution_order.append("middleware_after")
# Create ChatAgent with middleware
middleware = StreamingTrackingMiddleware()
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Set up mock streaming responses
mock_chat_client.streaming_responses = [
[
ChatResponseUpdate(contents=[TextContent(text="Streaming")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[TextContent(text=" response")], role=Role.ASSISTANT),
]
]
# Execute streaming
messages = [ChatMessage(role=Role.USER, text="test message")]
updates: list[AgentRunResponseUpdate] = []
async for update in agent.run_stream(messages):
updates.append(update)
# Verify streaming response
assert len(updates) == 2
assert updates[0].text == "Streaming"
assert updates[1].text == " response"
assert mock_chat_client.call_count == 1
# Verify middleware was called and streaming flag was set correctly
assert execution_order == ["middleware_before", "middleware_after"]
assert streaming_flags == [True] # Context should indicate streaming
async def test_non_streaming_vs_streaming_flag_validation(self, mock_chat_client: MockChatClient) -> None:
"""Test that is_streaming flag is correctly set for different execution modes."""
streaming_flags: list[bool] = []
class FlagTrackingMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
streaming_flags.append(context.is_streaming)
await next(context)
# Create ChatAgent with middleware
middleware = FlagTrackingMiddleware()
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
messages = [ChatMessage(role=Role.USER, text="test message")]
# Test non-streaming execution
response = await agent.run(messages)
assert response is not None
# Test streaming execution
async for _ in agent.run_stream(messages):
pass
# Verify flags: [non-streaming, streaming]
assert streaming_flags == [False, True]
class TestChatAgentMultipleMiddlewareOrdering:
"""Test cases for multiple middleware execution order with ChatAgent."""
async def test_multiple_agent_middleware_execution_order(self, mock_chat_client: MockChatClient) -> None:
"""Test that multiple agent middlewares execute in correct order with ChatAgent."""
execution_order: list[str] = []
class OrderedMiddleware(AgentMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Create multiple middlewares
middleware1 = OrderedMiddleware("first")
middleware2 = OrderedMiddleware("second")
middleware3 = OrderedMiddleware("third")
# Create ChatAgent with multiple middlewares
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware1, middleware2, middleware3])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert mock_chat_client.call_count == 1
# Verify execution order (should be nested: first wraps second wraps third)
expected_order = ["first_before", "second_before", "third_before", "third_after", "second_after", "first_after"]
assert execution_order == expected_order
async def test_mixed_middleware_types_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
"""Test mixed class and function-based middlewares with ChatAgent."""
execution_order: list[str] = []
class ClassAgentMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("class_agent_before")
await next(context)
execution_order.append("class_agent_after")
async def function_agent_middleware(
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("function_agent_before")
await next(context)
execution_order.append("function_agent_after")
class ClassFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("class_function_before")
await next(context)
execution_order.append("class_function_after")
async def function_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
execution_order.append("function_function_after")
# Create ChatAgent with mixed middleware types (no tools, focusing on agent middleware)
agent = ChatAgent(
chat_client=mock_chat_client,
middleware=[
ClassAgentMiddleware(),
function_agent_middleware,
ClassFunctionMiddleware(), # Won't execute without function calls
function_function_middleware, # Won't execute without function calls
],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert mock_chat_client.call_count == 1
# Verify that agent middlewares were executed in correct order
# (Function middlewares won't execute since no functions are called)
expected_order = ["class_agent_before", "function_agent_before", "function_agent_after", "class_agent_after"]
assert execution_order == expected_order
# region Tool Functions for Testing
def sample_tool_function(location: str) -> str:
"""A simple tool function for middleware testing."""
return f"Weather in {location}: sunny"
# region ChatAgent Function Middleware Tests with Tools
class TestChatAgentFunctionMiddlewareWithTools:
"""Test cases for function middleware integration with ChatAgent when tools are used."""
async def test_class_based_function_middleware_with_tool_calls(self, mock_chat_client: MockChatClient) -> None:
"""Test class-based function middleware with ChatAgent when function calls are made."""
execution_order: list[str] = []
class TrackingFunctionMiddleware(FunctionMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id="call_123",
name="sample_tool_function",
arguments='{"location": "Seattle"}',
)
],
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
mock_chat_client.responses = [function_call_response, final_response]
# Create ChatAgent with function middleware and tools
middleware = TrackingFunctionMiddleware("function_middleware")
agent = ChatAgent(
chat_client=mock_chat_client,
middleware=[middleware],
tools=[sample_tool_function],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for Seattle")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert mock_chat_client.call_count == 2 # Two calls: one for function call, one for final response
# Verify function middleware was executed
assert execution_order == ["function_middleware_before", "function_middleware_after"]
# Verify function call and result are in the response
all_contents = [content for message in response.messages for content in message.contents]
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
assert len(function_calls) == 1
assert len(function_results) == 1
assert function_calls[0].name == "sample_tool_function"
assert function_results[0].call_id == function_calls[0].call_id
async def test_function_based_function_middleware_with_tool_calls(self, mock_chat_client: MockChatClient) -> None:
"""Test function-based function middleware with ChatAgent when function calls are made."""
execution_order: list[str] = []
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_middleware_before")
await next(context)
execution_order.append("function_middleware_after")
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id="call_456",
name="sample_tool_function",
arguments='{"location": "San Francisco"}',
)
],
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
mock_chat_client.responses = [function_call_response, final_response]
# Create ChatAgent with function middleware and tools
agent = ChatAgent(
chat_client=mock_chat_client,
middleware=[tracking_function_middleware],
tools=[sample_tool_function],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for San Francisco")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert mock_chat_client.call_count == 2 # Two calls: one for function call, one for final response
# Verify function middleware was executed
assert execution_order == ["function_middleware_before", "function_middleware_after"]
# Verify function call and result are in the response
all_contents = [content for message in response.messages for content in message.contents]
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
assert len(function_calls) == 1
assert len(function_results) == 1
assert function_calls[0].name == "sample_tool_function"
assert function_results[0].call_id == function_calls[0].call_id
async def test_mixed_agent_and_function_middleware_with_tool_calls(self, mock_chat_client: MockChatClient) -> None:
"""Test both agent and function middleware with ChatAgent when function calls are made."""
execution_order: list[str] = []
class TrackingAgentMiddleware(AgentMiddleware):
async def process(
self,
context: AgentRunContext,
next: Callable[[AgentRunContext], Awaitable[None]],
) -> None:
execution_order.append("agent_middleware_before")
await next(context)
execution_order.append("agent_middleware_after")
class TrackingFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("function_middleware_before")
await next(context)
execution_order.append("function_middleware_after")
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id="call_789",
name="sample_tool_function",
arguments='{"location": "New York"}',
)
],
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
mock_chat_client.responses = [function_call_response, final_response]
# Create ChatAgent with both agent and function middleware and tools
agent = ChatAgent(
chat_client=mock_chat_client,
middleware=[TrackingAgentMiddleware(), TrackingFunctionMiddleware()],
tools=[sample_tool_function],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for New York")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert mock_chat_client.call_count == 2 # Two calls: one for function call, one for final response
# Verify middleware execution order: agent middleware wraps everything,
# function middleware only for function calls
expected_order = [
"agent_middleware_before",
"function_middleware_before",
"function_middleware_after",
"agent_middleware_after",
]
assert execution_order == expected_order
# Verify function call and result are in the response
all_contents = [content for message in response.messages for content in message.contents]
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
assert len(function_calls) == 1
assert len(function_results) == 1
assert function_calls[0].name == "sample_tool_function"
assert function_results[0].call_id == function_calls[0].call_id
python/packages/main/tests/main/test_middleware_context_result.py
+463
View File
@@ -0,0 +1,463 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import AsyncIterable, Awaitable, Callable
from typing import Any
from unittest.mock import MagicMock
import pytest
from pydantic import BaseModel, Field
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentRunResponseUpdate,
ChatAgent,
ChatMessage,
Role,
TextContent,
)
from agent_framework._middleware import (
AgentMiddleware,
AgentMiddlewarePipeline,
AgentRunContext,
FunctionInvocationContext,
FunctionMiddleware,
FunctionMiddlewarePipeline,
)
from agent_framework._tools import AIFunction
from .conftest import MockChatClient
class FunctionTestArgs(BaseModel):
"""Test arguments for function middleware tests."""
name: str = Field(description="Test name parameter")
class TestResultOverrideMiddleware:
"""Test cases for middleware result override functionality."""
async def test_agent_middleware_response_override_non_streaming(self, mock_agent: AgentProtocol) -> None:
"""Test that agent middleware can override response for non-streaming execution."""
override_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="overridden response")])
class ResponseOverrideMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Execute the pipeline first, then override the response
await next(context)
context.response = override_response
middleware = ResponseOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
handler_called = False
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
nonlocal handler_called
handler_called = True
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="original response")])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
# Verify the overridden response is returned
assert result is not None
assert result == override_response
assert result.messages[0].text == "overridden response"
# Verify original handler was called since middleware called next()
assert handler_called
async def test_agent_middleware_response_override_streaming(self, mock_agent: AgentProtocol) -> None:
"""Test that agent middleware can override response for streaming execution."""
async def override_stream() -> AsyncIterable[AgentRunResponseUpdate]:
yield AgentRunResponseUpdate(contents=[TextContent(text="overridden")])
yield AgentRunResponseUpdate(contents=[TextContent(text=" stream")])
class StreamResponseOverrideMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Execute the pipeline first, then override the response stream
await next(context)
context.response = override_stream()
middleware = StreamResponseOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
yield AgentRunResponseUpdate(contents=[TextContent(text="original")])
updates: list[AgentRunResponseUpdate] = []
async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
updates.append(update)
# Verify the overridden response stream is returned
assert len(updates) == 2
assert updates[0].text == "overridden"
assert updates[1].text == " stream"
async def test_function_middleware_result_override(self, mock_function: AIFunction[Any, Any]) -> None:
"""Test that function middleware can override result."""
override_result = "overridden function result"
class ResultOverrideMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Execute the pipeline first, then override the result
await next(context)
context.result = override_result
middleware = ResultOverrideMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
arguments = FunctionTestArgs(name="test")
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
handler_called = False
async def final_handler(ctx: FunctionInvocationContext) -> str:
nonlocal handler_called
handler_called = True
return "original function result"
result = await pipeline.execute(mock_function, arguments, context, final_handler)
# Verify the overridden result is returned
assert result == override_result
# Verify original handler was called since middleware called next()
assert handler_called
async def test_chat_agent_middleware_response_override(self) -> None:
"""Test result override functionality with ChatAgent integration."""
mock_chat_client = MockChatClient()
class ChatAgentResponseOverrideMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Always call next() first to allow execution
await next(context)
# Then conditionally override based on content
if any("special" in msg.text for msg in context.messages if msg.text):
context.response = AgentRunResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="Special response from middleware!")]
)
# Create ChatAgent with override middleware
middleware = ChatAgentResponseOverrideMiddleware()
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Test override case
override_messages = [ChatMessage(role=Role.USER, text="Give me a special response")]
override_response = await agent.run(override_messages)
assert override_response.messages[0].text == "Special response from middleware!"
# Verify chat client was called since middleware called next()
assert mock_chat_client.call_count == 1
# Test normal case
normal_messages = [ChatMessage(role=Role.USER, text="Normal request")]
normal_response = await agent.run(normal_messages)
assert normal_response.messages[0].text == "test response"
# Verify chat client was called for normal case
assert mock_chat_client.call_count == 2
async def test_chat_agent_middleware_streaming_override(self) -> None:
"""Test streaming result override functionality with ChatAgent integration."""
mock_chat_client = MockChatClient()
async def custom_stream() -> AsyncIterable[AgentRunResponseUpdate]:
yield AgentRunResponseUpdate(contents=[TextContent(text="Custom")])
yield AgentRunResponseUpdate(contents=[TextContent(text=" streaming")])
yield AgentRunResponseUpdate(contents=[TextContent(text=" response!")])
class ChatAgentStreamOverrideMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Always call next() first to allow execution
await next(context)
# Then conditionally override based on content
if any("custom stream" in msg.text for msg in context.messages if msg.text):
context.response = custom_stream()
# Create ChatAgent with override middleware
middleware = ChatAgentStreamOverrideMiddleware()
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
# Test streaming override case
override_messages = [ChatMessage(role=Role.USER, text="Give me a custom stream")]
override_updates: list[AgentRunResponseUpdate] = []
async for update in agent.run_stream(override_messages):
override_updates.append(update)
assert len(override_updates) == 3
assert override_updates[0].text == "Custom"
assert override_updates[1].text == " streaming"
assert override_updates[2].text == " response!"
# Test normal streaming case
normal_messages = [ChatMessage(role=Role.USER, text="Normal streaming request")]
normal_updates: list[AgentRunResponseUpdate] = []
async for update in agent.run_stream(normal_messages):
normal_updates.append(update)
assert len(normal_updates) == 2
assert normal_updates[0].text == "test streaming response "
assert normal_updates[1].text == "another update"
async def test_agent_middleware_conditional_no_next(self, mock_agent: AgentProtocol) -> None:
"""Test that when agent middleware conditionally doesn't call next(), no execution happens."""
class ConditionalNoNextMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Only call next() if message contains "execute"
if any("execute" in msg.text for msg in context.messages if msg.text):
await next(context)
# Otherwise, don't call next() - no execution should happen
middleware = ConditionalNoNextMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
handler_called = False
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
nonlocal handler_called
handler_called = True
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
# Test case where next() is NOT called
no_execute_messages = [ChatMessage(role=Role.USER, text="Don't run this")]
no_execute_context = AgentRunContext(agent=mock_agent, messages=no_execute_messages)
no_execute_result = await pipeline.execute(mock_agent, no_execute_messages, no_execute_context, final_handler)
# When middleware doesn't call next(), result should be empty AgentRunResponse
assert no_execute_result is not None
assert isinstance(no_execute_result, AgentRunResponse)
assert no_execute_result.messages == [] # Empty response
assert not handler_called
assert no_execute_context.response is None
# Reset for next test
handler_called = False
# Test case where next() IS called
execute_messages = [ChatMessage(role=Role.USER, text="Please execute this")]
execute_context = AgentRunContext(agent=mock_agent, messages=execute_messages)
execute_result = await pipeline.execute(mock_agent, execute_messages, execute_context, final_handler)
assert execute_result is not None
assert execute_result.messages[0].text == "executed response"
assert handler_called
async def test_function_middleware_conditional_no_next(self, mock_function: AIFunction[Any, Any]) -> None:
"""Test that when function middleware conditionally doesn't call next(), no execution happens."""
class ConditionalNoNextFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Only call next() if argument name contains "execute"
args = context.arguments
assert isinstance(args, FunctionTestArgs)
if "execute" in args.name:
await next(context)
# Otherwise, don't call next() - no execution should happen
middleware = ConditionalNoNextFunctionMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
handler_called = False
async def final_handler(ctx: FunctionInvocationContext) -> str:
nonlocal handler_called
handler_called = True
return "executed function result"
# Test case where next() is NOT called
no_execute_args = FunctionTestArgs(name="test_no_action")
no_execute_context = FunctionInvocationContext(function=mock_function, arguments=no_execute_args)
no_execute_result = await pipeline.execute(mock_function, no_execute_args, no_execute_context, final_handler)
# When middleware doesn't call next(), function result should be None (functions can return None)
assert no_execute_result is None
assert not handler_called
assert no_execute_context.result is None
# Reset for next test
handler_called = False
# Test case where next() IS called
execute_args = FunctionTestArgs(name="test_execute")
execute_context = FunctionInvocationContext(function=mock_function, arguments=execute_args)
execute_result = await pipeline.execute(mock_function, execute_args, execute_context, final_handler)
assert execute_result == "executed function result"
assert handler_called
class TestResultObservability:
"""Test cases for middleware result observability functionality."""
async def test_agent_middleware_response_observability(self, mock_agent: AgentProtocol) -> None:
"""Test that middleware can observe response after execution."""
observed_responses: list[AgentRunResponse] = []
class ObservabilityMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Context should be empty before next()
assert context.response is None
# Call next to execute
await next(context)
# Context should now contain the response for observability
assert context.response is not None
assert isinstance(context.response, AgentRunResponse)
observed_responses.append(context.response)
middleware = ObservabilityMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
# Verify response was observed
assert len(observed_responses) == 1
assert observed_responses[0].messages[0].text == "executed response"
assert result == observed_responses[0]
async def test_function_middleware_result_observability(self, mock_function: AIFunction[Any, Any]) -> None:
"""Test that middleware can observe function result after execution."""
observed_results: list[str] = []
class ObservabilityMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Context should be empty before next()
assert context.result is None
# Call next to execute
await next(context)
# Context should now contain the result for observability
assert context.result is not None
observed_results.append(context.result)
middleware = ObservabilityMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
arguments = FunctionTestArgs(name="test")
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
async def final_handler(ctx: FunctionInvocationContext) -> str:
return "executed function result"
result = await pipeline.execute(mock_function, arguments, context, final_handler)
# Verify result was observed
assert len(observed_results) == 1
assert observed_results[0] == "executed function result"
assert result == observed_results[0]
async def test_agent_middleware_post_execution_override(self, mock_agent: AgentProtocol) -> None:
"""Test that middleware can override response after observing execution."""
class PostExecutionOverrideMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
# Call next to execute first
await next(context)
# Now observe and conditionally override
assert context.response is not None
assert isinstance(context.response, AgentRunResponse)
if "modify" in context.response.messages[0].text:
# Override after observing
context.response = AgentRunResponse(
messages=[ChatMessage(role=Role.ASSISTANT, text="modified after execution")]
)
middleware = PostExecutionOverrideMiddleware()
pipeline = AgentMiddlewarePipeline([middleware])
messages = [ChatMessage(role=Role.USER, text="test")]
context = AgentRunContext(agent=mock_agent, messages=messages)
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response to modify")])
result = await pipeline.execute(mock_agent, messages, context, final_handler)
# Verify response was modified after execution
assert result is not None
assert result.messages[0].text == "modified after execution"
async def test_function_middleware_post_execution_override(self, mock_function: AIFunction[Any, Any]) -> None:
"""Test that middleware can override function result after observing execution."""
class PostExecutionOverrideMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Call next to execute first
await next(context)
# Now observe and conditionally override
assert context.result is not None
if "modify" in context.result:
# Override after observing
context.result = "modified after execution"
middleware = PostExecutionOverrideMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
arguments = FunctionTestArgs(name="test")
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
async def final_handler(ctx: FunctionInvocationContext) -> str:
return "result to modify"
result = await pipeline.execute(mock_function, arguments, context, final_handler)
# Verify result was modified after execution
assert result == "modified after execution"
@pytest.fixture
def mock_agent() -> AgentProtocol:
"""Mock agent for testing."""
agent = MagicMock(spec=AgentProtocol)
agent.name = "test_agent"
return agent
@pytest.fixture
def mock_function() -> AIFunction[Any, Any]:
"""Mock function for testing."""
function = MagicMock(spec=AIFunction[Any, Any])
function.name = "test_function"
return function
python/packages/main/tests/main/test_middleware_with_agent.py
+544
View File
@@ -0,0 +1,544 @@
# Copyright (c) Microsoft. All rights reserved.
from collections.abc import Awaitable, Callable
from agent_framework import (
AgentRunResponseUpdate,
ChatAgent,
ChatMessage,
ChatResponse,
ChatResponseUpdate,
FunctionCallContent,
FunctionResultContent,
Role,
TextContent,
)
from agent_framework._middleware import (
AgentMiddleware,
AgentRunContext,
FunctionInvocationContext,
FunctionMiddleware,
)
from .conftest import MockChatClient
# region ChatAgent Tests
class TestChatAgentClassBasedMiddleware:
"""Test cases for class-based middleware integration with ChatAgent."""
async def test_class_based_agent_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
"""Test class-based agent middleware with ChatAgent."""
execution_order: list[str] = []
class TrackingAgentMiddleware(AgentMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Create ChatAgent with middleware
middleware = TrackingAgentMiddleware("agent_middleware")
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
# Note: conftest "MockChatClient" returns different text format
assert "test response" in response.messages[0].text
# Verify middleware execution order
assert execution_order == ["agent_middleware_before", "agent_middleware_after"]
async def test_class_based_function_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
"""Test class-based function middleware with ChatAgent."""
execution_order: list[str] = []
class TrackingFunctionMiddleware(FunctionMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
middleware = TrackingFunctionMiddleware("function_middleware")
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert chat_client.call_count == 1
# Note: Function middleware won't execute since no function calls are made
assert execution_order == []
class TestChatAgentFunctionBasedMiddleware:
"""Test cases for function-based middleware integration with ChatAgent."""
async def test_function_based_agent_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
"""Test function-based agent middleware with ChatAgent."""
execution_order: list[str] = []
async def tracking_agent_middleware(
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("agent_function_before")
await next(context)
execution_order.append("agent_function_after")
# Create ChatAgent with function middleware
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_agent_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert response.messages[0].role == Role.ASSISTANT
assert response.messages[0].text == "test response"
assert chat_client.call_count == 1
# Verify middleware execution order
assert execution_order == ["agent_function_before", "agent_function_after"]
async def test_function_based_function_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
"""Test function-based function middleware with ChatAgent."""
execution_order: list[str] = []
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
execution_order.append("function_function_after")
# Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_function_middleware])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert chat_client.call_count == 1
# Note: Function middleware won't execute since no function calls are made
assert execution_order == []
class TestChatAgentStreamingMiddleware:
"""Test cases for streaming middleware integration with ChatAgent."""
async def test_agent_middleware_with_streaming(self, chat_client: "MockChatClient") -> None:
"""Test agent middleware with streaming ChatAgent responses."""
execution_order: list[str] = []
streaming_flags: list[bool] = []
class StreamingTrackingMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("middleware_before")
streaming_flags.append(context.is_streaming)
await next(context)
execution_order.append("middleware_after")
# Create ChatAgent with middleware
middleware = StreamingTrackingMiddleware()
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
# Set up mock streaming responses
chat_client.streaming_responses = [
[
ChatResponseUpdate(contents=[TextContent(text="Streaming")], role=Role.ASSISTANT),
ChatResponseUpdate(contents=[TextContent(text=" response")], role=Role.ASSISTANT),
]
]
# Execute streaming
messages = [ChatMessage(role=Role.USER, text="test message")]
updates: list[AgentRunResponseUpdate] = []
async for update in agent.run_stream(messages):
updates.append(update)
# Verify streaming response
assert len(updates) == 2
assert updates[0].text == "Streaming"
assert updates[1].text == " response"
assert chat_client.call_count == 1
# Verify middleware was called and streaming flag was set correctly
assert execution_order == ["middleware_before", "middleware_after"]
assert streaming_flags == [True] # Context should indicate streaming
async def test_non_streaming_vs_streaming_flag_validation(self, chat_client: "MockChatClient") -> None:
"""Test that is_streaming flag is correctly set for different execution modes."""
streaming_flags: list[bool] = []
class FlagTrackingMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
streaming_flags.append(context.is_streaming)
await next(context)
# Create ChatAgent with middleware
middleware = FlagTrackingMiddleware()
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
messages = [ChatMessage(role=Role.USER, text="test message")]
# Test non-streaming execution
response = await agent.run(messages)
assert response is not None
# Test streaming execution
async for _ in agent.run_stream(messages):
pass
# Verify flags: [non-streaming, streaming]
assert streaming_flags == [False, True]
class TestChatAgentMultipleMiddlewareOrdering:
"""Test cases for multiple middleware execution order with ChatAgent."""
async def test_multiple_agent_middleware_execution_order(self, chat_client: "MockChatClient") -> None:
"""Test that multiple agent middlewares execute in correct order with ChatAgent."""
execution_order: list[str] = []
class OrderedMiddleware(AgentMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Create multiple middlewares
middleware1 = OrderedMiddleware("first")
middleware2 = OrderedMiddleware("second")
middleware3 = OrderedMiddleware("third")
# Create ChatAgent with multiple middlewares
agent = ChatAgent(chat_client=chat_client, middleware=[middleware1, middleware2, middleware3])
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert chat_client.call_count == 1
# Verify execution order (should be nested: first wraps second wraps third)
expected_order = ["first_before", "second_before", "third_before", "third_after", "second_after", "first_after"]
assert execution_order == expected_order
async def test_mixed_middleware_types_with_chat_agent(self, chat_client: "MockChatClient") -> None:
"""Test mixed class and function-based middlewares with ChatAgent."""
execution_order: list[str] = []
class ClassAgentMiddleware(AgentMiddleware):
async def process(
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("class_agent_before")
await next(context)
execution_order.append("class_agent_after")
async def function_agent_middleware(
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
) -> None:
execution_order.append("function_agent_before")
await next(context)
execution_order.append("function_agent_after")
class ClassFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("class_function_before")
await next(context)
execution_order.append("class_function_after")
async def function_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
execution_order.append("function_function_after")
# Create ChatAgent with mixed middleware types (no tools, focusing on agent middleware)
agent = ChatAgent(
chat_client=chat_client,
middleware=[
ClassAgentMiddleware(),
function_agent_middleware,
ClassFunctionMiddleware(), # Won't execute without function calls
function_function_middleware, # Won't execute without function calls
],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="test message")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert chat_client.call_count == 1
# Verify that agent middlewares were executed in correct order
# (Function middlewares won't execute since no functions are called)
expected_order = ["class_agent_before", "function_agent_before", "function_agent_after", "class_agent_after"]
assert execution_order == expected_order
# region Tool Functions for Testing
def sample_tool_function(location: str) -> str:
"""A simple tool function for middleware testing."""
return f"Weather in {location}: sunny"
# region ChatAgent Function Middleware Tests with Tools
class TestChatAgentFunctionMiddlewareWithTools:
"""Test cases for function middleware integration with ChatAgent when tools are used."""
async def test_class_based_function_middleware_with_tool_calls(self, chat_client: "MockChatClient") -> None:
"""Test class-based function middleware with ChatAgent when function calls are made."""
execution_order: list[str] = []
class TrackingFunctionMiddleware(FunctionMiddleware):
def __init__(self, name: str):
self.name = name
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
execution_order.append(f"{self.name}_after")
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id="call_123",
name="sample_tool_function",
arguments='{"location": "Seattle"}',
)
],
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
chat_client.responses = [function_call_response, final_response]
# Create ChatAgent with function middleware and tools
middleware = TrackingFunctionMiddleware("function_middleware")
agent = ChatAgent(
chat_client=chat_client,
middleware=[middleware],
tools=[sample_tool_function],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for Seattle")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert chat_client.call_count == 2 # Two calls: one for function call, one for final response
# Verify function middleware was executed
assert execution_order == ["function_middleware_before", "function_middleware_after"]
# Verify function call and result are in the response
all_contents = [content for message in response.messages for content in message.contents]
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
assert len(function_calls) == 1
assert len(function_results) == 1
assert function_calls[0].name == "sample_tool_function"
assert function_results[0].call_id == function_calls[0].call_id
async def test_function_based_function_middleware_with_tool_calls(self, chat_client: "MockChatClient") -> None:
"""Test function-based function middleware with ChatAgent when function calls are made."""
execution_order: list[str] = []
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_middleware_before")
await next(context)
execution_order.append("function_middleware_after")
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id="call_456",
name="sample_tool_function",
arguments='{"location": "San Francisco"}',
)
],
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
chat_client.responses = [function_call_response, final_response]
# Create ChatAgent with function middleware and tools
agent = ChatAgent(
chat_client=chat_client,
middleware=[tracking_function_middleware],
tools=[sample_tool_function],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for San Francisco")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert chat_client.call_count == 2 # Two calls: one for function call, one for final response
# Verify function middleware was executed
assert execution_order == ["function_middleware_before", "function_middleware_after"]
# Verify function call and result are in the response
all_contents = [content for message in response.messages for content in message.contents]
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
assert len(function_calls) == 1
assert len(function_results) == 1
assert function_calls[0].name == "sample_tool_function"
assert function_results[0].call_id == function_calls[0].call_id
async def test_mixed_agent_and_function_middleware_with_tool_calls(self, chat_client: "MockChatClient") -> None:
"""Test both agent and function middleware with ChatAgent when function calls are made."""
execution_order: list[str] = []
class TrackingAgentMiddleware(AgentMiddleware):
async def process(
self,
context: AgentRunContext,
next: Callable[[AgentRunContext], Awaitable[None]],
) -> None:
execution_order.append("agent_middleware_before")
await next(context)
execution_order.append("agent_middleware_after")
class TrackingFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("function_middleware_before")
await next(context)
execution_order.append("function_middleware_after")
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
ChatMessage(
role=Role.ASSISTANT,
contents=[
FunctionCallContent(
call_id="call_789",
name="sample_tool_function",
arguments='{"location": "New York"}',
)
],
)
]
)
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
chat_client.responses = [function_call_response, final_response]
# Create ChatAgent with both agent and function middleware and tools
agent = ChatAgent(
chat_client=chat_client,
middleware=[TrackingAgentMiddleware(), TrackingFunctionMiddleware()],
tools=[sample_tool_function],
)
# Execute the agent
messages = [ChatMessage(role=Role.USER, text="Get weather for New York")]
response = await agent.run(messages)
# Verify response
assert response is not None
assert len(response.messages) > 0
assert chat_client.call_count == 2 # Two calls: one for function call, one for final response
# Verify middleware execution order: agent middleware wraps everything,
# function middleware only for function calls
expected_order = [
"agent_middleware_before",
"function_middleware_before",
"function_middleware_after",
"agent_middleware_after",
]
assert execution_order == expected_order
# Verify function call and result are in the response
all_contents = [content for message in response.messages for content in message.contents]
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
assert len(function_calls) == 1
assert len(function_results) == 1
assert function_calls[0].name == "sample_tool_function"
assert function_results[0].call_id == function_calls[0].call_id