TestSystemMessageUpdateViaMiddleware Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/test_system_message.py lines 362–449

class TestSystemMessageUpdateViaMiddleware:
    """Test updating system messages through middleware."""

    def test_middleware_can_set_initial_system_message(self) -> None:
        """Test middleware setting system message when none exists."""

        def set_system_message_middleware(
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelResponse:
            """Middleware that sets initial system message."""
            new_request = request.override(
                system_message=SystemMessage(content="Set by middleware")
            )
            return handler(new_request)

        model = GenericFakeChatModel(messages=iter([AIMessage(content="response")]))
        request = ModelRequest(
            model=model,
            system_message=None,
            messages=[HumanMessage(content="Hello")],
            tool_choice=None,
            tools=[],
            response_format=None,
            state=AgentState(messages=[]),
            runtime=Runtime(),
        )

        captured_request = None

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

        set_system_message_middleware(request, mock_handler)

        assert captured_request is not None
        assert captured_request.system_message is not None
        assert len(captured_request.system_message.content_blocks) == 1
        assert captured_request.system_message.content_blocks[0].get("text") == "Set by middleware"

    def test_middleware_can_update_via_system_message_object(self) -> None:
        """Test middleware updating system message using SystemMessage objects."""

        def append_with_metadata_middleware(
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelResponse:
            """Append using SystemMessage to preserve metadata."""
            base_content = request.system_message.text if request.system_message else ""
            base_kwargs = request.system_message.additional_kwargs if request.system_message else {}

            new_message = SystemMessage(
                content=base_content + " Additional instructions.",
                additional_kwargs={**base_kwargs, "middleware": "applied"},
            )
            new_request = request.override(system_message=new_message)
            return handler(new_request)

        model = GenericFakeChatModel(messages=iter([AIMessage(content="response")]))
        request = ModelRequest(
            model=model,
            system_message=SystemMessage(
                content="Base prompt", additional_kwargs={"base": "value"}
            ),
            messages=[],
            tool_choice=None,
            tools=[],
            response_format=None,
            state=AgentState(messages=[]),
            runtime=Runtime(),
        )

        captured_request = None

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