test_multiple_transformations() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 840–873

    def test_multiple_transformations(self) -> None:
        """Test multiple middleware that each transform the response."""

        class PrefixMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                result = handler(request)
                # result is ModelResponse, extract AIMessage from it
                ai_message = result.result[0]
                return AIMessage(content=f"[PREFIX] {ai_message.content}")

        class SuffixMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                result = handler(request)
                # result is ModelResponse, extract AIMessage from it
                ai_message = result.result[0]
                return AIMessage(content=f"{ai_message.content} [SUFFIX]")

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Middle")]))
        # Prefix is outer, Suffix is inner
        # Inner (Suffix) runs first, then Outer (Prefix)
        agent = create_agent(model=model, middleware=[PrefixMiddleware(), SuffixMiddleware()])

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

        # Suffix adds suffix first, then Prefix adds prefix
        assert result["messages"][1].content == "[PREFIX] Middle [SUFFIX]"