test_todo_middleware_on_model_call() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_todo.py lines 136–171

def test_todo_middleware_on_model_call(
    original_prompt: str | None, expected_prompt_prefix: str
) -> None:
    """Test that wrap_model_call handles system prompts correctly."""
    middleware = TodoListMiddleware()
    model = FakeToolCallingModel()

    state: PlanningState = {"messages": [HumanMessage(content="Hello")]}

    request = ModelRequest(
        model=model,
        system_prompt=original_prompt,
        messages=[HumanMessage(content="Hello")],
        tool_choice=None,
        tools=[],
        response_format=None,
        state=state,
        runtime=cast("Runtime", object()),
        model_settings={},
    )

    captured_request = None

    def mock_handler(req: ModelRequest) -> ModelResponse:
        nonlocal captured_request
        captured_request = req
        return ModelResponse(result=[AIMessage(content="mock response")])

    # Call wrap_model_call to trigger the middleware logic
    middleware.wrap_model_call(request, mock_handler)
    # Check that the modified request passed to handler has the expected prompt
    assert captured_request is not None
    assert captured_request.system_prompt is not None
    assert captured_request.system_prompt.startswith(expected_prompt_prefix)
    # Original request should be unchanged
    assert request.system_prompt == original_prompt