diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 0372967ccf..45050dfb1f 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -2,7 +2,6 @@ from __future__ import annotations -import contextlib import inspect import logging import sys @@ -1452,8 +1451,13 @@ async def current_handler() -> None: return current_handler first_handler = create_next_handler(0) - with contextlib.suppress(MiddlewareTermination): + try: await first_handler() + except MiddlewareTermination as exc: + if context.result is None: + context.result = cast( + "AgentResponse | ResponseStream[AgentResponseUpdate, AgentResponse] | None", exc.result + ) if context.result and isinstance(context.result, ResponseStream): if context.stream_buffer_updates: @@ -1663,8 +1667,13 @@ async def current_handler() -> None: return current_handler first_handler = create_next_handler(0) - with contextlib.suppress(MiddlewareTermination): + try: await first_handler() + except MiddlewareTermination as exc: + if context.result is None: + context.result = cast( + "ChatResponse | ResponseStream[ChatResponseUpdate, ChatResponse] | None", exc.result + ) if context.result and isinstance(context.result, ResponseStream): if context.stream_buffer_updates: diff --git a/python/packages/core/tests/core/test_middleware.py b/python/packages/core/tests/core/test_middleware.py index ab473306b5..b74012c7f7 100644 --- a/python/packages/core/tests/core/test_middleware.py +++ b/python/packages/core/tests/core/test_middleware.py @@ -32,6 +32,8 @@ ) from agent_framework._tools import FunctionTool, ToolTypes +from .utils import termination_result + class TestAgentContext: """Test cases for AgentContext.""" @@ -1929,3 +1931,65 @@ def test_as_middleware_list_owns_the_bare_source_rule(self) -> None: assert _as_middleware_list((agent_mw,)) == [agent_mw] # Strings are sequences but never element-ized into characters. assert _as_middleware_list("bare-string") == ["bare-string"] # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + +class TestMiddlewareTerminationResult: + """``MiddlewareTermination(result=...)`` is the run's result in the agent and chat pipelines too.""" + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("kind", ["agent", "chat"]) + async def test_termination_result_is_returned( + self, kind: str, stream: bool, mock_agent: SupportsAgentRun, mock_chat_client: Any + ) -> None: + expected = termination_result(kind, stream) + handler_calls: list[str] = [] + + async def terminate(context: Any, call_next: Callable[[], Awaitable[None]]) -> None: + raise MiddlewareTermination("blocked", result=expected) + + async def final_handler(context: Any) -> Any: + handler_calls.append("final") + return None + + messages = [Message(role="user", contents=["test"])] + result: object + if kind == "agent": + result = await AgentMiddlewarePipeline(terminate).execute( + AgentContext(agent=mock_agent, messages=messages, stream=stream), final_handler + ) + else: + result = await ChatMiddlewarePipeline(terminate).execute( + ChatContext(client=mock_chat_client, messages=messages, options={}, stream=stream), final_handler + ) + + assert result is expected + assert handler_calls == [] + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("kind", ["agent", "chat"]) + async def test_context_result_wins_over_termination_result( + self, kind: str, stream: bool, mock_agent: SupportsAgentRun, mock_chat_client: Any + ) -> None: + """A result the middleware already set on the context wins, as in the function pipeline.""" + from_context = termination_result(kind, stream) + from_exception = termination_result(kind, stream) + + async def set_then_terminate(context: Any, call_next: Callable[[], Awaitable[None]]) -> None: + context.result = from_context + raise MiddlewareTermination("blocked", result=from_exception) + + async def final_handler(context: Any) -> Any: + return None + + messages = [Message(role="user", contents=["test"])] + result: object + if kind == "agent": + result = await AgentMiddlewarePipeline(set_then_terminate).execute( + AgentContext(agent=mock_agent, messages=messages, stream=stream), final_handler + ) + else: + result = await ChatMiddlewarePipeline(set_then_terminate).execute( + ChatContext(client=mock_chat_client, messages=messages, options={}, stream=stream), final_handler + ) + + assert result is from_context diff --git a/python/packages/core/tests/core/test_middleware_with_agent.py b/python/packages/core/tests/core/test_middleware_with_agent.py index d04767c904..9b06e69833 100644 --- a/python/packages/core/tests/core/test_middleware_with_agent.py +++ b/python/packages/core/tests/core/test_middleware_with_agent.py @@ -37,6 +37,7 @@ from agent_framework._sessions import InMemoryHistoryProvider from .conftest import MockBaseChatClient, MockChatClient +from .utils import TERMINATION_TEXT, termination_result # region Agent Tests @@ -3275,3 +3276,34 @@ async def process(self, context: AgentContext, call_next: Callable[[], Awaitable # endregion + + +class TestMiddlewareTerminationResultWithAgent: + """A middleware that terminates with ``result=`` answers ``Agent.run`` without calling the model.""" + + @pytest.mark.parametrize("stream", [False, True], ids=["non_streaming", "streaming"]) + @pytest.mark.parametrize("kind", ["agent", "chat"]) + async def test_agent_run_returns_termination_result( + self, kind: str, stream: bool, chat_client_base: MockBaseChatClient + ) -> None: + @agent_middleware + async def block_agent(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: + raise MiddlewareTermination("blocked", result=termination_result("agent", context.stream)) + + @chat_middleware + async def block_chat(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None: + raise MiddlewareTermination("blocked", result=termination_result("chat", context.stream)) + + agent = Agent(client=chat_client_base, middleware=[block_agent if kind == "agent" else block_chat]) + + if stream: + response_stream = agent.run("hello", stream=True) + updates = [update async for update in response_stream] + response = await response_stream.get_final_response() + assert "".join(update.text for update in updates) == TERMINATION_TEXT + else: + response = await agent.run("hello") + + assert response is not None + assert response.text == TERMINATION_TEXT + assert chat_client_base.call_count == 0 diff --git a/python/packages/core/tests/core/utils.py b/python/packages/core/tests/core/utils.py index 61bf277d94..ad32fbfeb6 100644 --- a/python/packages/core/tests/core/utils.py +++ b/python/packages/core/tests/core/utils.py @@ -1,9 +1,42 @@ # Copyright (c) Microsoft. All rights reserved. +from collections.abc import AsyncIterable from copy import deepcopy +from typing import Any from unittest.mock import MagicMock +from agent_framework import ( + AgentResponse, + AgentResponseUpdate, + ChatResponse, + ChatResponseUpdate, + Content, + Message, + ResponseStream, +) + +TERMINATION_TEXT = "Sorry, blocked." + + +def termination_result(kind: str, stream: bool) -> Any: + """The response a middleware hands to ``MiddlewareTermination(result=...)``.""" + if kind == "agent": + if not stream: + return AgentResponse(messages=[Message(role="assistant", contents=[TERMINATION_TEXT])]) + + async def _agent_updates() -> AsyncIterable[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(TERMINATION_TEXT)], role="assistant") + + return ResponseStream(_agent_updates(), finalizer=AgentResponse.from_updates) + if not stream: + return ChatResponse(messages=[Message(role="assistant", contents=[TERMINATION_TEXT])]) + + async def _chat_updates() -> AsyncIterable[ChatResponseUpdate]: + yield ChatResponseUpdate(contents=[Content.from_text(TERMINATION_TEXT)], role="assistant") + + return ResponseStream(_chat_updates(), finalizer=ChatResponse.from_updates) + class CopyingMock(MagicMock): def __call__(self, *args, **kwargs):