test_decorator_returns_wrap_result() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call_state_update.py lines 580–608

    def test_decorator_returns_wrap_result(self) -> None:
        """@wrap_model_call decorator can return ExtendedModelResponse with command."""

        @wrap_model_call
        def command_middleware(
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ExtendedModelResponse:
            response = handler(request)
            return ExtendedModelResponse(
                model_response=response,
                command=Command(
                    update={
                        "messages": [
                            HumanMessage(content="From decorator", id="dec"),
                        ]
                    }
                ),
            )

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Model response")]))
        agent = create_agent(model=model, middleware=[command_middleware])

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

        messages = result["messages"]
        assert len(messages) == 3
        assert messages[1].content == "Model response"
        assert messages[2].content == "From decorator"