diff --git a/headroom/backends/litellm.py b/headroom/backends/litellm.py index 526a8d76e..36987b3e5 100644 --- a/headroom/backends/litellm.py +++ b/headroom/backends/litellm.py @@ -747,6 +747,39 @@ class LiteLLMBackend(Backend): return converted + def _system_field_to_message(self, system: Any) -> dict[str, Any]: + """Convert Anthropic's top-level `system` field to an OpenAI-style message. + + `system` can be a plain string or a list of content blocks, each of + which may carry its own `cache_control` breakpoint (Claude Code puts + the prompt-caching marker on the last system block). Flattening the + list to a joined string, as this code used to do, drops that + `cache_control` entirely: litellm's Bedrock Converse transformation + only emits a `cachePoint` when it sees content blocks with + `cache_control` on them, never for a plain string. That silently + broke prompt caching of the system prefix. #1390 covers the analogous + case for tool_result blocks in `_convert_messages_for_litellm` above; + this handles the top-level `system` field, which was out of scope + there. Preserve block structure and cache_control so the breakpoint + survives into the litellm call. + """ + if isinstance(system, str): + return {"role": "system", "content": system} + if isinstance(system, list): + blocks: list[dict[str, Any]] = [] + for s in system: + if isinstance(s, dict): + block: dict[str, Any] = {"type": "text", "text": s.get("text", "")} + if "cache_control" in s: + block["cache_control"] = s["cache_control"] + else: + block = {"type": "text", "text": str(s)} + blocks.append(block) + return {"role": "system", "content": blocks} + # Shouldn't happen in practice (None is filtered out via "system" in + # body), but stay defensive rather than raising. + return {"role": "system", "content": str(system)} + def _to_anthropic_response( self, litellm_response: Any, @@ -843,15 +876,7 @@ class LiteLLMBackend(Backend): # System prompt (Anthropic puts it in body, OpenAI in messages) if "system" in body: - system = body["system"] - if isinstance(system, str): - kwargs["messages"].insert(0, {"role": "system", "content": system}) - elif isinstance(system, list): - # Anthropic list format - system_text = " ".join( - s.get("text", "") if isinstance(s, dict) else str(s) for s in system - ) - kwargs["messages"].insert(0, {"role": "system", "content": system_text}) + kwargs["messages"].insert(0, self._system_field_to_message(body["system"])) # Provider-specific region config if self.region: @@ -956,14 +981,7 @@ class LiteLLMBackend(Backend): if "tool_choice" in body: kwargs["tool_choice"] = _convert_tool_choice(body["tool_choice"]) if "system" in body: - system = body["system"] - if isinstance(system, str): - kwargs["messages"].insert(0, {"role": "system", "content": system}) - elif isinstance(system, list): - system_text = " ".join( - s.get("text", "") if isinstance(s, dict) else str(s) for s in system - ) - kwargs["messages"].insert(0, {"role": "system", "content": system_text}) + kwargs["messages"].insert(0, self._system_field_to_message(body["system"])) # Provider-specific region config if self.region: diff --git a/tests/test_bedrock_tool_result_cache_and_streaming_stats.py b/tests/test_bedrock_tool_result_cache_and_streaming_stats.py index f29e0ccc7..59591f011 100644 --- a/tests/test_bedrock_tool_result_cache_and_streaming_stats.py +++ b/tests/test_bedrock_tool_result_cache_and_streaming_stats.py @@ -252,3 +252,69 @@ class TestStreamingCacheStatsCompletion: assert usage["input_tokens"] == 42 assert "cache_read_input_tokens" not in usage assert "cache_creation_input_tokens" not in usage + + +class TestSystemFieldCacheControl: + """The top-level Anthropic `system` field must not be flattened to a + plain string when it carries per-block cache_control, or Bedrock prompt + caching of the system prefix silently breaks (see module docstring, + #1390's uncovered case).""" + + def test_list_system_with_cache_control_preserved(self): + backend = _backend() + system = [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Long static prefix.", "cache_control": {"type": "ephemeral"}}, + ] + msg = backend._system_field_to_message(system) + assert msg["role"] == "system" + assert isinstance(msg["content"], list) + assert msg["content"][-1]["cache_control"] == {"type": "ephemeral"} + + def test_string_system_unaffected(self): + backend = _backend() + msg = backend._system_field_to_message("You are a helpful assistant.") + assert msg == {"role": "system", "content": "You are a helpful assistant."} + assert isinstance(msg["content"], str) + + def test_list_system_without_cache_control_has_no_cache_control_keys(self): + backend = _backend() + system = [ + {"type": "text", "text": "First block."}, + {"type": "text", "text": "Second block."}, + ] + msg = backend._system_field_to_message(system) + assert isinstance(msg["content"], list) + assert all("cache_control" not in block for block in msg["content"]) + + def test_bedrock_converse_transform_emits_cachepoint_for_list_with_cache_control(self): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + backend = _backend() + system = [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Long static prefix.", "cache_control": {"type": "ephemeral"}}, + ] + system_msg = backend._system_field_to_message(system) + messages = [system_msg, {"role": "user", "content": "hi"}] + + _, system_blocks = AmazonConverseConfig()._transform_system_message( + messages, model="global.anthropic.claude-sonnet-5" + ) + assert any("cachePoint" in block for block in system_blocks) + + def test_bedrock_converse_transform_omits_cachepoint_without_cache_control(self): + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + backend = _backend() + system = [ + {"type": "text", "text": "First block."}, + {"type": "text", "text": "Second block."}, + ] + system_msg = backend._system_field_to_message(system) + messages = [system_msg, {"role": "user", "content": "hi"}] + + _, system_blocks = AmazonConverseConfig()._transform_system_message( + messages, model="global.anthropic.claude-sonnet-5" + ) + assert not any("cachePoint" in block for block in system_blocks)