test_on_model_call_decorator() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_decorators.py lines 98–133

def test_on_model_call_decorator() -> None:
    """Test wrap_model_call decorator with all configuration options."""

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

    # Verify all options were applied
    assert isinstance(custom_on_model_call, AgentMiddleware)
    assert custom_on_model_call.state_schema == CustomState
    assert custom_on_model_call.tools == [test_tool]
    assert custom_on_model_call.__class__.__name__ == "CustomOnModelCall"

    # Verify it works
    original_request = ModelRequest(
        model=FakeToolCallingModel(),
        system_prompt="Original",
        messages=[HumanMessage("Hello")],
        tool_choice=None,
        tools=[],
        response_format=None,
        state={"messages": [HumanMessage("Hello")]},
        runtime=None,
    )

    def mock_handler(req: ModelRequest) -> ModelResponse:
        return ModelResponse(
            result=[AIMessage(content=f"Handled with prompt: {req.system_prompt}")]
        )

    result = custom_on_model_call.wrap_model_call(original_request, mock_handler)
    assert isinstance(result, ModelResponse)
    assert result.result[0].content == "Handled with prompt: Modified"