test_three_middleware_composition() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 720–774

    def test_three_middleware_composition(self) -> None:
        """Test composition of three middleware."""
        execution_order = []

        class FirstMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                execution_order.append("first-before")
                response = handler(request)
                execution_order.append("first-after")
                return response

        class SecondMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                execution_order.append("second-before")
                response = handler(request)
                execution_order.append("second-after")
                return response

        class ThirdMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                execution_order.append("third-before")
                response = handler(request)
                execution_order.append("third-after")
                return response

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Response")]))
        agent = create_agent(
            model=model,
            middleware=[FirstMiddleware(), SecondMiddleware(), ThirdMiddleware()],
        )

        agent.invoke({"messages": [HumanMessage("Test")]})

        # First wraps Second wraps Third:
        # 1-before, 2-before, 3-before, model, 3-after, 2-after, 1-after
        assert execution_order == [
            "first-before",
            "second-before",
            "third-before",
            "third-after",
            "second-after",
            "first-after",
        ]