test_async_model_with_middleware() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 1382–1403

    async def test_async_model_with_middleware(self) -> None:
        """Test that wrap_model_call works with async model execution."""
        log = []

        class LoggingMiddleware(AgentMiddleware):
            async def awrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
            ) -> ModelCallResult:
                log.append("before")
                result = await handler(request)
                log.append("after")
                return result

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

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

        assert log == ["before", "after"]
        assert result["messages"][1].content == "Async response"