test_wrap_model_call_state_update.py — langchain Source File

Source Code

"""Unit tests for ExtendedModelResponse command support in wrap_model_call.

Tests that wrap_model_call middleware can return ExtendedModelResponse to provide
a Command alongside the model response. Commands are applied as separate state
updates through graph reducers (e.g. add_messages for messages).
"""

from collections.abc import Awaitable, Callable

import pytest
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.errors import InvalidUpdateError
from langgraph.types import Command

from langchain.agents import AgentState, create_agent
from langchain.agents.middleware.types import (
    AgentMiddleware,
    ExtendedModelResponse,
    ModelRequest,
    ModelResponse,
    wrap_model_call,
)


class TestBasicCommand:
    """Test basic ExtendedModelResponse functionality with Command."""

    def test_command_messages_added_alongside_model_messages(self) -> None:
        """Command messages are added alongside model response messages (additive)."""

        class AddMessagesMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ExtendedModelResponse:
                response = handler(request)
                custom_msg = HumanMessage(content="Custom message", id="custom")
                return ExtendedModelResponse(
                    model_response=response,
                    command=Command(update={"messages": [custom_msg]}),
                )

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Hello!")]))
        agent = create_agent(model=model, middleware=[AddMessagesMiddleware()])

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

        # Both model response AND command messages appear (additive via add_messages)
        messages = result["messages"]
        assert len(messages) == 3
        assert messages[0].content == "Hi"
        assert messages[1].content == "Hello!"
        assert messages[2].content == "Custom message"

    def test_command_with_extra_messages_and_model_response(self) -> None:
        """Middleware can add extra messages via command alongside model messages."""

        class ExtraMessagesMiddleware(AgentMiddleware):
// ... (858 more lines)