test_cache_control_preserved_across_middleware() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/test_system_message.py lines 624–679

    def test_cache_control_preserved_across_middleware(self) -> None:
        """Test that cache control is preserved when middleware modifies message."""

        def first_middleware_with_cache(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            """Set system message with cache control."""
            new_message = SystemMessage(
                content=[
                    {
                        "type": "text",
                        "text": "Cached content",
                        "cache_control": {"type": "ephemeral"},
                    }
                ]
            )
            new_request = request.override(system_message=new_message)
            return handler(new_request)

        def second_middleware_appends(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            """Append to system message while preserving cache control."""
            assert request.system_message is not None
            existing_content = request.system_message.content_blocks
            new_content = [*existing_content, TextContentBlock(type="text", text="Additional text")]

            new_message = SystemMessage(content_blocks=new_content)
            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=None,
            messages=[],
            tool_choice=None,
            tools=[],
            response_format=None,
            state=AgentState(messages=[]),
            runtime=Runtime(),
        )

        def final_handler(req: ModelRequest) -> ModelResponse:
            # Verify cache control was preserved
            assert req.system_message is not None
            assert isinstance(req.system_message.content_blocks, list)
            assert len(req.system_message.content_blocks) == 2
            assert req.system_message.content_blocks[0].get("cache_control") == {
                "type": "ephemeral"
            }
            return ModelResponse(result=[AIMessage(content="response")])

        first_middleware_with_cache(
            request, lambda req: second_middleware_appends(req, final_handler)
        )