mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(backend/bedrock): preserve system-prompt cache_control breakpoint (list form) (#2225)
## Description The LiteLLM backend flattened the Anthropic top-level `system` field to a joined string whenever it arrived as a **list of content blocks**, discarding each block's `cache_control`. LiteLLM's Bedrock Converse transform (`AmazonConverseConfig._transform_system_message`) only emits a `cachePoint` for content blocks that carry `cache_control`, never for a plain string. So on any `--backend bedrock` deployment the **system prefix was never cached**: every turn re-sent the full system prompt (typically 5k-25k tokens with Claude Code) at full input price. #1390 fixed the analogous case for `tool_result` blocks in `_convert_messages_for_litellm`, but the top-level `system` field handling in `send_message` / `stream_message` was out of scope there and still flattened. The cache hits observed on live Bedrock traffic came only from the tool-result / message-tail breakpoint, masking that the largest, most stable block was uncached. Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/backends/litellm.py`: factor the top-level `system` field conversion into a single `_system_field_to_message` helper. `str` stays string-content (unchanged behavior); a `list` maps to text blocks retaining each block's `cache_control`; non-dict entries coerce to a plain text block. Both call sites (`send_message` non-streaming, `stream_message` streaming) now call the helper, so they stay byte-identical. - `tests/test_bedrock_tool_result_cache_and_streaming_stats.py`: add `TestSystemFieldCacheControl` — list-with-`cache_control` retains it, plain-string is unchanged, list-without-`cache_control` produces list content with no marker, plus two end-to-end checks that drive the emitted message through `AmazonConverseConfig._transform_system_message` and assert a `cachePoint` is present for the cache_control case and absent otherwise. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ uv run pytest tests/test_bedrock_tool_result_cache_and_streaming_stats.py -q collected 13 items tests/test_bedrock_tool_result_cache_and_streaming_stats.py ............. [100%] 13 passed in 1.18s $ uv run ruff check headroom/backends/litellm.py tests/test_bedrock_tool_result_cache_and_streaming_stats.py All checks passed! ``` ## Real Behavior Proof - Environment: personal fork deployed as a real proxy (macOS launchd service, `headroom install apply`) with `--backend bedrock --mode cache`, fronting a live Claude Code session. Model `global.anthropic.claude-sonnet-5`, region eu-west-1. - Exact command / steps: ran a purpose-built probe that POSTs Anthropic-shape `/v1/messages` to the running proxy with a 7,692-token STABLE system prompt carrying a single `cache_control: {type: ephemeral}` breakpoint (and no other cache_control anywhere), a pinned `x-headroom-session-id`, across 5 sequential turns, reading the raw response `usage` each turn. - Observed result: **before the fix**, the response `usage` had no cache fields at all — `cache_creation_input_tokens` and `cache_read_input_tokens` both absent, nothing cached. **After the fix**, turn 1 shows `cache_creation_input_tokens=10164` (write) and turns 2-5 each show `cache_read_input_tokens=10164` (read) with `cache_creation=0` — write-once, then read the system prefix from Bedrock's cache on every subsequent turn. The proxy's `/stats` `prefix_cache` tracker registered all four later turns as hits (`hit_requests += 1` per turn, `bust_count = 0`). - Not tested: no change to the tool_result / message-tail breakpoint path (already handled by #1390 / #2144); this fix is scoped to the top-level `system` field only. The in-`messages` text-block flatten in `_convert_messages_for_litellm` is intentionally left untouched. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes - No linked issue number: found via independent investigation of a personal `--backend bedrock` deployment. - Companion to #2196 (`fix(proxy/bedrock): wire PrefixCacheTracker updates into Bedrock backend paths`) from the same investigation. #2196 wires the tracker; this fixes the system-prompt breakpoint that #1390 left flattened on the top-level `system` field. Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
842d7e1ad1
commit
ea0115cbdb
2 changed files with 101 additions and 17 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue