Python: [BREAKING] Refactor middleware layering and split Anthropic raw client (#4746)

* [BREAKING] Refactor middleware layering and raw clients Reorder chat client layers so function invocation wraps chat middleware, and chat middleware stays outside telemetry while still running for each inner model call. Add middleware pipeline caching, refresh docs and samples, and split Anthropic into raw and public clients to match the standard layering model. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Tighten typing ignores in ancillary modules Add targeted typing ignores in workflow visualization and lab modules so pyright stays clean alongside the middleware refactor work. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix categorize_middleware to unpack tuple/Sequence and use relative MRO assertions - Broaden isinstance check in categorize_middleware from list to Sequence so tuples and other Sequence types are properly unpacked instead of being appended as a single item. - Replace fragile hardcoded MRO index assertions in anthropic test with relative ordering via mro.index(). - Add regression tests for categorize_middleware with tuple, list, and None inputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix middleware string decomposition, add middleware param to FunctionInvocationLayer, and add tests (#4710) - Guard categorize_middleware Sequence check against str/bytes to prevent character-by-character decomposition of accidentally passed strings - Add explicit middleware parameter to FunctionInvocationLayer.get_response and merge it into client_kwargs before categorization, fixing the inconsistency where only OpenAIChatClient supported this parameter - Add assertions that RawAnthropicClient does not inherit convenience layers - Add chat middleware cache test with non-empty base middleware - Add tests for single unwrapped middleware item and string input Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Apply pre-commit auto-fixes * Apply pre-commit auto-fixes * Address review feedback for #4710: review comment fixes --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <copilot@github.com>
2026-06-16 21:04:09 +08:00 · 2026-03-20 01:43:37 +01:00
parent cefda44283
commit 0cd40f8354
41 changed files with 936 additions and 155 deletions
@@ -128,8 +128,8 @@ class MockChatClient:


 class MockBaseChatClient(
-    ChatMiddlewareLayer[OptionsCoT],
    FunctionInvocationLayer[OptionsCoT],
+    ChatMiddlewareLayer[OptionsCoT],
    ChatTelemetryLayer[OptionsCoT],
    BaseChatClient[OptionsCoT],
    Generic[OptionsCoT],
@@ -137,7 +137,7 @@ class MockBaseChatClient(
    """Mock implementation of a full-featured ChatClient."""

    def __init__(self, **kwargs: Any):
-        super().__init__(function_middleware=[], **kwargs)
+        super().__init__(middleware=[], **kwargs)
        self.run_responses: list[ChatResponse] = []
        self.streaming_responses: list[list[ChatResponseUpdate]] = []
        self.call_count: int = 0
@@ -74,8 +74,8 @@ def test_openai_chat_client_get_response_docstring_surfaces_layered_runtime_docs
    assert docstring is not None
    assert "Get a response from a chat client." in docstring
    assert "function_invocation_kwargs" in docstring
-    assert "function_middleware: Optional per-call function middleware." in docstring
    assert "middleware: Optional per-call chat and function middleware." in docstring
+    assert "function_middleware: Optional per-call function middleware." not in docstring


 def test_openai_chat_client_get_response_is_defined_on_openai_class() -> None:
@@ -84,7 +84,6 @@ def test_openai_chat_client_get_response_is_defined_on_openai_class() -> None:
    signature = inspect.signature(OpenAIChatClient.get_response)

    assert OpenAIChatClient.get_response.__qualname__ == "OpenAIChatClient.get_response"
-    assert "function_middleware" in signature.parameters
    assert "middleware" in signature.parameters


@@ -3226,7 +3226,7 @@ async def test_terminate_loop_single_function_call(chat_client_base: SupportsCha
    response = await chat_client_base.get_response(
        "hello",
        options={"tool_choice": "auto", "tools": [ai_func]},
-        middleware=[TerminateLoopMiddleware()],
+        client_kwargs={"middleware": [TerminateLoopMiddleware()]},
    )

    # Function should NOT have been executed - middleware intercepted it
@@ -3292,7 +3292,7 @@ async def test_terminate_loop_multiple_function_calls_one_terminates(chat_client
    response = await chat_client_base.get_response(
        "hello",
        options={"tool_choice": "auto", "tools": [normal_func, terminating_func]},
-        middleware=[SelectiveTerminateMiddleware()],
+        client_kwargs={"middleware": [SelectiveTerminateMiddleware()]},
    )

    # normal_function should have executed (middleware calls next_handler)
@@ -3345,7 +3345,7 @@ async def test_terminate_loop_streaming_single_function_call(chat_client_base: S
    async for update in chat_client_base.get_response(
        "hello",
        options={"tool_choice": "auto", "tools": [ai_func]},
-        middleware=[TerminateLoopMiddleware()],
+        client_kwargs={"middleware": [TerminateLoopMiddleware()]},
        stream=True,
    ):
        updates.append(update)
@@ -3389,12 +3389,12 @@ async def test_conversation_id_updated_in_options_between_tool_iterations():
    conversation_ids_received: list[str | None] = []

    class TrackingChatClient(
-        ChatMiddlewareLayer,
        FunctionInvocationLayer,
+        ChatMiddlewareLayer,
        BaseChatClient,
    ):
        def __init__(self) -> None:
-            super().__init__(function_middleware=[])
+            super().__init__(middleware=[])
            self.run_responses: list[ChatResponse] = []
            self.streaming_responses: list[list[ChatResponseUpdate]] = []
            self.call_count: int = 0
@@ -84,8 +84,8 @@ class _MockBaseChatClient(BaseChatClient[Any]):


 class FunctionInvokingMockClient(
-    ChatMiddlewareLayer[Any],
    FunctionInvocationLayer[Any],
+    ChatMiddlewareLayer[Any],
    ChatTelemetryLayer[Any],
    _MockBaseChatClient,
 ):
@@ -28,6 +28,7 @@ from agent_framework._middleware import (
    FunctionMiddleware,
    FunctionMiddlewarePipeline,
    MiddlewareTermination,
+    categorize_middleware,
 )
 from agent_framework._tools import FunctionTool

@@ -1681,3 +1682,49 @@ def mock_chat_client() -> Any:
    client = MagicMock(spec=SupportsChatGetResponse)
    client.service_url = MagicMock(return_value="mock://test")
    return client
+
+
+class TestCategorizeMiddleware:
+    """Test cases for categorize_middleware."""
+
+    def test_categorize_middleware_with_tuple(self) -> None:
+        """Test that tuple middleware sources are unpacked, not appended as a single item."""
+        chat_mw = TestChatMiddleware()
+        function_mw = TestFunctionMiddleware()
+        agent_mw = TestAgentMiddleware()
+        result = categorize_middleware((chat_mw, function_mw, agent_mw))
+        assert result["chat"] == [chat_mw]
+        assert result["function"] == [function_mw]
+        assert result["agent"] == [agent_mw]
+
+    def test_categorize_middleware_with_list(self) -> None:
+        """Test that list middleware sources are unpacked correctly."""
+        chat_mw = TestChatMiddleware()
+        function_mw = TestFunctionMiddleware()
+        result = categorize_middleware([chat_mw, function_mw])
+        assert result["chat"] == [chat_mw]
+        assert result["function"] == [function_mw]
+        assert result["agent"] == []
+
+    def test_categorize_middleware_with_none(self) -> None:
+        """Test that None middleware sources are handled."""
+        result = categorize_middleware(None)
+        assert result["chat"] == []
+        assert result["function"] == []
+        assert result["agent"] == []
+
+    def test_categorize_middleware_with_single_item(self) -> None:
+        """Test that a single unwrapped middleware item is appended correctly."""
+        chat_mw = TestChatMiddleware()
+        result = categorize_middleware(chat_mw)
+        assert result["chat"] == [chat_mw]
+        assert result["function"] == []
+        assert result["agent"] == []
+
+    def test_categorize_middleware_with_string_does_not_decompose(self) -> None:
+        """Test that a string is not decomposed character-by-character."""
+        result = categorize_middleware("not_a_middleware")
+        # String should be treated as a single item, not decomposed into characters
+        total_items = len(result["chat"]) + len(result["function"]) + len(result["agent"])
+        assert total_items == 1
+        assert result["agent"] == ["not_a_middleware"]
@@ -697,6 +697,26 @@ class TestChatAgentFunctionMiddlewareWithTools:
        assert function_calls[0].name == "sample_tool_function"
        assert function_results[0].call_id == function_calls[0].call_id

+    def test_agent_middleware_pipeline_cache_reuses_matching_middleware(self) -> None:
+        """Test that identical agent middleware sets reuse the cached pipeline."""
+
+        @agent_middleware
+        async def first_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        @agent_middleware
+        async def second_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        agent = Agent(client=MockBaseChatClient())
+
+        first_pipeline = agent._get_agent_middleware_pipeline([first_middleware])
+        second_pipeline = agent._get_agent_middleware_pipeline([first_middleware])
+        third_pipeline = agent._get_agent_middleware_pipeline([second_middleware])
+
+        assert first_pipeline is second_pipeline
+        assert third_pipeline is not first_pipeline
+
    async def test_function_middleware_can_access_and_override_custom_kwargs(
        self, chat_client_base: "MockBaseChatClient"
    ) -> None:
@@ -1969,6 +1989,77 @@ class TestChatAgentChatMiddleware:
            "agent_middleware_after",
        ]

+    async def test_combined_middleware_with_tool_loop(self) -> None:
+        """Test Agent middleware ordering when tool calls trigger multiple chat rounds."""
+        execution_order: list[str] = []
+        chat_round = 0
+        client = MockBaseChatClient()
+        client.run_responses = [
+            ChatResponse(
+                messages=[
+                    Message(
+                        role="assistant",
+                        contents=[
+                            Content.from_function_call(
+                                call_id="call_123",
+                                name="sample_tool_function",
+                                arguments='{"location": "Seattle"}',
+                            )
+                        ],
+                    )
+                ]
+            ),
+            ChatResponse(messages=[Message(role="assistant", text="Final response")]),
+        ]
+
+        async def tracking_agent_middleware(
+            context: AgentContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            execution_order.append("agent_middleware_before")
+            await call_next()
+            execution_order.append("agent_middleware_after")
+
+        async def tracking_chat_middleware(
+            context: ChatContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            nonlocal chat_round
+            chat_round += 1
+            execution_order.append(f"chat_middleware_before_{chat_round}")
+            await call_next()
+            execution_order.append(f"chat_middleware_after_{chat_round}")
+
+        async def tracking_function_middleware(
+            context: FunctionInvocationContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            execution_order.append("function_middleware_before")
+            await call_next()
+            execution_order.append("function_middleware_after")
+
+        agent = Agent(
+            client=client,
+            middleware=[tracking_chat_middleware, tracking_function_middleware, tracking_agent_middleware],
+            tools=[sample_tool_function],
+        )
+
+        response = await agent.run([Message(role="user", text="test")])
+
+        assert response is not None
+        assert client.call_count == 2
+        assert response.messages[-1].text == "Final response"
+        assert execution_order == [
+            "agent_middleware_before",
+            "chat_middleware_before_1",
+            "chat_middleware_after_1",
+            "function_middleware_before",
+            "function_middleware_after",
+            "chat_middleware_before_2",
+            "chat_middleware_after_2",
+            "agent_middleware_after",
+        ]
+
    async def test_agent_middleware_can_access_and_override_custom_kwargs(self) -> None:
        """Test that agent middleware can access and override custom parameters like temperature."""
        captured_kwargs: dict[str, Any] = {}
@@ -274,7 +274,10 @@ class TestChatMiddleware:

        # First call with run-level middleware
        messages = [Message(role="user", text="first message")]
-        response1 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
+        response1 = await chat_client_base.get_response(
+            messages,
+            client_kwargs={"middleware": [counting_middleware]},
+        )
        assert response1 is not None
        assert execution_count["count"] == 1

@@ -286,7 +289,10 @@ class TestChatMiddleware:

        # Third call with run-level middleware again - should execute
        messages = [Message(role="user", text="third message")]
-        response3 = await chat_client_base.get_response(messages, middleware=[counting_middleware])
+        response3 = await chat_client_base.get_response(
+            messages,
+            client_kwargs={"middleware": [counting_middleware]},
+        )
        assert response3 is not None
        assert execution_count["count"] == 2  # Should be 2 now

@@ -335,6 +341,81 @@ class TestChatMiddleware:
        assert modified_kwargs["new_param"] == "added_by_middleware"
        assert modified_kwargs["custom_param"] == "test_value"  # Should still be there

+    def test_chat_middleware_pipeline_cache_reuses_matching_middleware(
+        self,
+        chat_client_base: "MockBaseChatClient",
+    ) -> None:
+        """Test that identical chat middleware sets reuse the cached pipeline."""
+
+        @chat_middleware
+        async def first_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        @chat_middleware
+        async def second_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        first_pipeline = chat_client_base._get_chat_middleware_pipeline([first_middleware])
+        second_pipeline = chat_client_base._get_chat_middleware_pipeline([first_middleware])
+        third_pipeline = chat_client_base._get_chat_middleware_pipeline([second_middleware])
+
+        assert first_pipeline is second_pipeline
+        assert third_pipeline is not first_pipeline
+
+    def test_chat_middleware_pipeline_cache_includes_base_middleware(
+        self,
+        chat_client_base: "MockBaseChatClient",
+    ) -> None:
+        """Test that chat middleware cache key includes base middleware to prevent incorrect reuse."""
+
+        @chat_middleware
+        async def base_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        @chat_middleware
+        async def runtime_middleware(context: ChatContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        # Without base middleware
+        pipeline_no_base = chat_client_base._get_chat_middleware_pipeline([runtime_middleware])
+
+        # With base middleware
+        chat_client_base.chat_middleware = [base_middleware]
+        pipeline_with_base = chat_client_base._get_chat_middleware_pipeline([runtime_middleware])
+
+        assert pipeline_with_base is not pipeline_no_base
+
+    def test_function_middleware_pipeline_cache_reuses_matching_middleware(
+        self,
+        chat_client_base: "MockBaseChatClient",
+    ) -> None:
+        """Test that identical function middleware sets reuse the cached pipeline."""
+
+        @function_middleware
+        async def base_middleware(context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None:
+            await call_next()
+
+        @function_middleware
+        async def first_runtime_middleware(
+            context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
+        ) -> None:
+            await call_next()
+
+        @function_middleware
+        async def second_runtime_middleware(
+            context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]
+        ) -> None:
+            await call_next()
+
+        chat_client_base.function_middleware = [base_middleware]
+
+        first_pipeline = chat_client_base._get_function_middleware_pipeline([first_runtime_middleware])
+        second_pipeline = chat_client_base._get_function_middleware_pipeline([first_runtime_middleware])
+        third_pipeline = chat_client_base._get_function_middleware_pipeline([second_runtime_middleware])
+
+        assert first_pipeline is second_pipeline
+        assert third_pipeline is not first_pipeline
+
    async def test_function_middleware_registration_on_chat_client(
        self, chat_client_base: "MockBaseChatClient"
    ) -> None:
@@ -450,7 +531,9 @@ class TestChatMiddleware:
        # Execute the chat client directly with run-level middleware and tools
        messages = [Message(role="user", text="What's the weather in New York?")]
        response = await client.get_response(
-            messages, options={"tools": [sample_tool_wrapped]}, middleware=[run_level_function_middleware]
+            messages,
+            options={"tools": [sample_tool_wrapped]},
+            client_kwargs={"middleware": [run_level_function_middleware]},
        )

        # Verify response
@@ -463,3 +546,156 @@ class TestChatMiddleware:
            "run_level_function_middleware_before",
            "run_level_function_middleware_after",
        ]
+
+    async def test_run_level_chat_and_function_middleware_split_per_function_loop_round(self) -> None:
+        """Test mixed run-level middleware is split so chat middleware runs per model call."""
+        execution_order: list[str] = []
+        chat_round = 0
+
+        @chat_middleware
+        async def run_level_chat_middleware(
+            context: ChatContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            nonlocal chat_round
+            chat_round += 1
+            execution_order.append(f"chat_middleware_before_{chat_round}")
+            await call_next()
+            execution_order.append(f"chat_middleware_after_{chat_round}")
+
+        @function_middleware
+        async def run_level_function_middleware(
+            context: FunctionInvocationContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            execution_order.append("function_middleware_before")
+            await call_next()
+            execution_order.append("function_middleware_after")
+
+        def sample_tool(location: str) -> str:
+            """Get weather for a location."""
+            return f"Weather in {location}: sunny"
+
+        sample_tool_wrapped = FunctionTool(
+            func=sample_tool,
+            name="sample_tool",
+            description="Get weather for a location",
+            approval_mode="never_require",
+        )
+
+        client = MockBaseChatClient()
+        client.run_responses = [
+            ChatResponse(
+                messages=[
+                    Message(
+                        role="assistant",
+                        contents=[
+                            Content.from_function_call(
+                                call_id="call_3",
+                                name="sample_tool",
+                                arguments={"location": "Seattle"},
+                            )
+                        ],
+                    )
+                ]
+            ),
+            ChatResponse(messages=[Message(role="assistant", text="Based on the weather data, it's sunny!")]),
+        ]
+
+        response = await client.get_response(
+            [Message(role="user", text="What's the weather in Seattle?")],
+            options={"tools": [sample_tool_wrapped]},
+            client_kwargs={"middleware": [run_level_chat_middleware, run_level_function_middleware]},
+        )
+
+        assert response is not None
+        assert client.call_count == 2
+        assert response.messages[-1].text == "Based on the weather data, it's sunny!"
+        assert execution_order == [
+            "chat_middleware_before_1",
+            "chat_middleware_after_1",
+            "function_middleware_before",
+            "function_middleware_after",
+            "chat_middleware_before_2",
+            "chat_middleware_after_2",
+        ]
+
+    async def test_run_level_chat_and_function_middleware_split_per_function_loop_round_streaming(self) -> None:
+        """Test mixed run-level middleware is split so chat middleware runs per model call in streaming mode."""
+        execution_order: list[str] = []
+        chat_round = 0
+
+        @chat_middleware
+        async def run_level_chat_middleware(
+            context: ChatContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            nonlocal chat_round
+            chat_round += 1
+            execution_order.append(f"chat_middleware_before_{chat_round}")
+            await call_next()
+            execution_order.append(f"chat_middleware_after_{chat_round}")
+
+        @function_middleware
+        async def run_level_function_middleware(
+            context: FunctionInvocationContext,
+            call_next: Callable[[], Awaitable[None]],
+        ) -> None:
+            execution_order.append("function_middleware_before")
+            await call_next()
+            execution_order.append("function_middleware_after")
+
+        def sample_tool(location: str) -> str:
+            """Get weather for a location."""
+            return f"Weather in {location}: sunny"
+
+        sample_tool_wrapped = FunctionTool(
+            func=sample_tool,
+            name="sample_tool",
+            description="Get weather for a location",
+            approval_mode="never_require",
+        )
+
+        client = MockBaseChatClient()
+        client.streaming_responses = [
+            [
+                ChatResponseUpdate(
+                    contents=[
+                        Content.from_function_call(
+                            call_id="call_3",
+                            name="sample_tool",
+                            arguments='{"location": "Seattle"}',
+                        )
+                    ],
+                    role="assistant",
+                    finish_reason="tool_calls",
+                ),
+            ],
+            [
+                ChatResponseUpdate(
+                    contents=[Content.from_text("Based on the weather data, it's sunny!")],
+                    role="assistant",
+                    finish_reason="stop",
+                ),
+            ],
+        ]
+
+        updates: list[ChatResponseUpdate] = []
+        async for update in client.get_response(
+            [Message(role="user", text="What's the weather in Seattle?")],
+            options={"tools": [sample_tool_wrapped]},
+            client_kwargs={"middleware": [run_level_chat_middleware, run_level_function_middleware]},
+            stream=True,
+        ):
+            updates.append(update)
+
+        assert client.call_count == 2
+        assert len(updates) > 0
+        assert execution_order == [
+            "chat_middleware_before_1",
+            "chat_middleware_after_1",
+            "function_middleware_before",
+            "function_middleware_after",
+            "chat_middleware_before_2",
+            "chat_middleware_after_2",
+        ]
@@ -2437,7 +2437,7 @@ def test_capture_response(span_exporter: InMemorySpanExporter):
 async def test_layer_ordering_span_sequence_with_function_calling(span_exporter: InMemorySpanExporter):
    """Test that with correct layer ordering, spans appear in the expected sequence.

-    When using the correct layer ordering (ChatMiddlewareLayer, FunctionInvocationLayer,
+    When using the correct layer ordering (FunctionInvocationLayer, ChatMiddlewareLayer,
    ChatTelemetryLayer, BaseChatClient), the spans should appear in this order:
    1. First 'chat' span (initial LLM call that returns function call)
    2. 'execute_tool' span (function invocation)
@@ -2454,11 +2454,11 @@ async def test_layer_ordering_span_sequence_with_function_calling(span_exporter:
    def get_weather(location: str) -> str:
        return f"The weather in {location} is sunny."

-    # Correct layer ordering: FunctionInvocationLayer BEFORE ChatTelemetryLayer
-    # This ensures each inner LLM call gets its own telemetry span
+    # Correct layer ordering: FunctionInvocationLayer BEFORE ChatMiddlewareLayer BEFORE ChatTelemetryLayer
+    # This ensures each inner LLM call traverses chat middleware and still gets its own telemetry span
    class MockChatClientWithLayers(
-        ChatMiddlewareLayer,
        FunctionInvocationLayer,
+        ChatMiddlewareLayer,
        ChatTelemetryLayer,
        BaseChatClient,
    ):