test_multi_stage_transformation() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 346–371

    def test_multi_stage_transformation(self) -> None:
        """Test middleware applying multiple transformations."""

        class MultiTransformMiddleware(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]

                # First transformation: uppercase
                assert isinstance(ai_message.content, str)
                content = ai_message.content.upper()
                # Second transformation: add prefix and suffix
                content = f"[START] {content} [END]"
                return AIMessage(content=content)

        model = GenericFakeChatModel(messages=iter([AIMessage(content="hello")]))
        agent = create_agent(model=model, middleware=[MultiTransformMiddleware()])

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

        assert result["messages"][1].content == "[START] HELLO [END]"