diff --git a/python/packages/main/agent_framework/_middleware.py b/python/packages/main/agent_framework/_middleware.py index 59f6dd5c29..9f7cac0bc4 100644 --- a/python/packages/main/agent_framework/_middleware.py +++ b/python/packages/main/agent_framework/_middleware.py @@ -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 diff --git a/python/packages/main/tests/main/conftest.py b/python/packages/main/tests/main/conftest.py index 001510f877..9cb938f8ea 100644 --- a/python/packages/main/tests/main/conftest.py +++ b/python/packages/main/tests/main/conftest.py @@ -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): diff --git a/python/packages/main/tests/main/test_middleware.py b/python/packages/main/tests/main/test_middleware.py index b49b473092..2d90064842 100644 --- a/python/packages/main/tests/main/test_middleware.py +++ b/python/packages/main/tests/main/test_middleware.py @@ -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 diff --git a/python/packages/main/tests/main/test_middleware_context_result.py b/python/packages/main/tests/main/test_middleware_context_result.py new file mode 100644 index 0000000000..86ca0e88ec --- /dev/null +++ b/python/packages/main/tests/main/test_middleware_context_result.py @@ -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 diff --git a/python/packages/main/tests/main/test_middleware_with_agent.py b/python/packages/main/tests/main/test_middleware_with_agent.py new file mode 100644 index 0000000000..7a2c030285 --- /dev/null +++ b/python/packages/main/tests/main/test_middleware_with_agent.py @@ -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