test_async_on_model_call_decorator() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_decorators.py lines 302–317

def test_async_on_model_call_decorator() -> None:
    """Test wrap_model_call decorator with async function."""

    @wrap_model_call(state_schema=CustomState, tools=[test_tool], name="AsyncOnModelCall")
    async def async_on_model_call(
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelCallResult:
        return await handler(
            request.override(system_message=SystemMessage(content="Modified async"))
        )

    assert isinstance(async_on_model_call, AgentMiddleware)
    assert async_on_model_call.state_schema == CustomState
    assert async_on_model_call.tools == [test_tool]
    assert async_on_model_call.__class__.__name__ == "AsyncOnModelCall"