From 4dccfcd9b7cf3905444b5b5a0160259289911451 Mon Sep 17 00:00:00 2001 From: joaomarcos Date: Thu, 23 Jul 2026 14:29:36 -0300 Subject: [PATCH] feat(bedrock): add Converse API prompt caching (cachePoint) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Claude-on-Bedrock already gets prompt caching via the AnthropicBedrock SDK path. This adds it to the raw Converse API path used for non-Claude models (Amazon Nova, and Claude when bearer-token auth forces Converse routing, #28156) — a conservative model allowlist inserts cachePoint blocks after tools, system, and the message before the newest turn, and extracts cacheReadInputTokens/cacheWriteInputTokens into usage so cost accounting picks them up through the existing Anthropic-style fallback. Ref: relatorio-cache-performance-provedores-ia.md, P0 item 1. --- agent/bedrock_adapter.py | 72 ++++++++++++++++++++---- tests/agent/test_bedrock_adapter.py | 86 +++++++++++++++++++++++++++++ tests/agent/test_usage_pricing.py | 21 +++++++ 3 files changed, 168 insertions(+), 11 deletions(-) diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py index c8cff3f76e1..c399081619f 100644 --- a/agent/bedrock_adapter.py +++ b/agent/bedrock_adapter.py @@ -433,6 +433,29 @@ def _model_supports_tool_use(model_id: str) -> bool: return not any(pattern in model_lower for pattern in _NON_TOOL_CALLING_PATTERNS) +# --------------------------------------------------------------------------- +# Prompt-cache capability detection (Converse API cachePoint) +# --------------------------------------------------------------------------- +# Claude on Bedrock already gets prompt caching through the AnthropicBedrock +# SDK path (see is_anthropic_bedrock_model / runtime_provider.py's dual-path +# routing) — it never reaches build_converse_kwargs unless bearer-token auth +# forces the Converse path (#28156). This allowlist covers the Converse API +# itself: sending an unsupported model a cachePoint block raises a +# ValidationException, so — like _model_supports_tool_use but inverted — +# unknown models default to NOT receiving cache markers until confirmed. +# Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html +_CACHE_POINT_PATTERNS = [ + "anthropic.claude", # bearer-token fallback path + "amazon.nova", +] + + +def _model_supports_prompt_cache(model_id: str) -> bool: + """Return True if the model accepts a Converse API cachePoint block.""" + model_lower = model_id.lower() + return any(pattern in model_lower for pattern in _CACHE_POINT_PATTERNS) + + def is_anthropic_bedrock_model(model_id: str) -> bool: """Return True if the model is an Anthropic Claude model on Bedrock. @@ -764,14 +787,22 @@ def normalize_converse_response(response: Dict) -> SimpleNamespace: reasoning_content="\n\n".join(reasoning_parts) if reasoning_parts else None, ) - # Build usage stats + # Build usage stats. Converse's inputTokens excludes cache read/write + # tokens (unlike OpenAI's prompt_tokens, which includes them) — restore + # the OpenAI-style "total includes cache" convention here so downstream + # normalize_usage() can subtract them back out consistently, and surface + # the Anthropic-named fields it already falls back to for cache reads. usage_data = response.get("usage", {}) + input_tokens = usage_data.get("inputTokens", 0) + cache_read_tokens = usage_data.get("cacheReadInputTokens", 0) + cache_write_tokens = usage_data.get("cacheWriteInputTokens", 0) + output_tokens = usage_data.get("outputTokens", 0) usage = SimpleNamespace( - prompt_tokens=usage_data.get("inputTokens", 0), - completion_tokens=usage_data.get("outputTokens", 0), - total_tokens=( - usage_data.get("inputTokens", 0) + usage_data.get("outputTokens", 0) - ), + prompt_tokens=input_tokens + cache_read_tokens + cache_write_tokens, + completion_tokens=output_tokens, + total_tokens=input_tokens + cache_read_tokens + cache_write_tokens + output_tokens, + cache_read_input_tokens=cache_read_tokens, + cache_creation_input_tokens=cache_write_tokens, ) finish_reason = _converse_stop_reason_to_openai(stop_reason) @@ -936,6 +967,8 @@ def stream_converse_with_callbacks( usage_data = { "inputTokens": meta_usage.get("inputTokens", 0), "outputTokens": meta_usage.get("outputTokens", 0), + "cacheReadInputTokens": meta_usage.get("cacheReadInputTokens", 0), + "cacheWriteInputTokens": meta_usage.get("cacheWriteInputTokens", 0), } # Flush remaining text @@ -949,12 +982,16 @@ def stream_converse_with_callbacks( reasoning_content="\n\n".join(reasoning_parts) if reasoning_parts else None, ) + input_tokens = usage_data.get("inputTokens", 0) + cache_read_tokens = usage_data.get("cacheReadInputTokens", 0) + cache_write_tokens = usage_data.get("cacheWriteInputTokens", 0) + output_tokens = usage_data.get("outputTokens", 0) usage = SimpleNamespace( - prompt_tokens=usage_data.get("inputTokens", 0), - completion_tokens=usage_data.get("outputTokens", 0), - total_tokens=( - usage_data.get("inputTokens", 0) + usage_data.get("outputTokens", 0) - ), + prompt_tokens=input_tokens + cache_read_tokens + cache_write_tokens, + completion_tokens=output_tokens, + total_tokens=input_tokens + cache_read_tokens + cache_write_tokens + output_tokens, + cache_read_input_tokens=cache_read_tokens, + cache_creation_input_tokens=cache_write_tokens, ) finish_reason = _converse_stop_reason_to_openai(stop_reason) @@ -993,6 +1030,7 @@ def build_converse_kwargs( Converts OpenAI-format inputs to Converse API parameters. """ system_prompt, converse_messages = convert_messages_to_converse(messages) + cache_enabled = _model_supports_prompt_cache(model) kwargs: Dict[str, Any] = { "modelId": model, @@ -1003,6 +1041,8 @@ def build_converse_kwargs( } if system_prompt: + if cache_enabled: + system_prompt = system_prompt + [{"cachePoint": {"type": "default"}}] kwargs["system"] = system_prompt from agent.anthropic_adapter import _forbids_sampling_params @@ -1026,6 +1066,8 @@ def build_converse_kwargs( # Strip tools for known non-tool-calling models and warn the user. # Ref: PR #7920 feedback from @ptlally, pattern from PR #4346. if _model_supports_tool_use(model): + if cache_enabled: + converse_tools = converse_tools + [{"cachePoint": {"type": "default"}}] kwargs["toolConfig"] = {"tools": converse_tools} else: logger.warning( @@ -1033,6 +1075,14 @@ def build_converse_kwargs( "The agent will operate in text-only mode.", model ) + if cache_enabled and len(converse_messages) >= 2: + # Checkpoint everything up to (not including) the newest turn, so the + # marker survives unchanged across requests as only the tail grows — + # mirroring the Anthropic system_and_3 strategy in prompt_caching.py. + content = converse_messages[-2].get("content") + if isinstance(content, list) and content: + content.append({"cachePoint": {"type": "default"}}) + if guardrail_config: kwargs["guardrailConfig"] = guardrail_config diff --git a/tests/agent/test_bedrock_adapter.py b/tests/agent/test_bedrock_adapter.py index 97c6d43689e..7226e4b7744 100644 --- a/tests/agent/test_bedrock_adapter.py +++ b/tests/agent/test_bedrock_adapter.py @@ -388,6 +388,29 @@ class TestNormalizeConverseResponse: assert result.usage.completion_tokens == 5 assert result.usage.total_tokens == 15 + def test_cache_tokens_folded_into_prompt_tokens(self): + """Converse's inputTokens excludes cache read/write tokens (unlike + OpenAI's prompt_tokens). normalize_converse_response must add them + back into prompt_tokens/total_tokens and surface the Anthropic-named + fields so normalize_usage() picks them up via its existing fallback.""" + from agent.bedrock_adapter import normalize_converse_response + response = { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 50, + "outputTokens": 20, + "cacheReadInputTokens": 900, + "cacheWriteInputTokens": 300, + }, + } + result = normalize_converse_response(response) + assert result.usage.prompt_tokens == 50 + 900 + 300 + assert result.usage.completion_tokens == 20 + assert result.usage.total_tokens == 50 + 900 + 300 + 20 + assert result.usage.cache_read_input_tokens == 900 + assert result.usage.cache_creation_input_tokens == 300 + def test_tool_use_response(self): from agent.bedrock_adapter import normalize_converse_response response = { @@ -700,6 +723,45 @@ class TestBuildConverseKwargs: ) assert "toolConfig" not in kwargs + def test_cache_point_added_for_supported_model(self): + """Claude and Nova on the Converse path get cachePoint markers on + system, tools, and the message before the newest turn.""" + from agent.bedrock_adapter import build_converse_kwargs + tools = [{"type": "function", "function": { + "name": "test", "description": "Test", "parameters": {}, + }}] + messages = [ + {"role": "system", "content": "Be helpful."}, + {"role": "user", "content": "First"}, + {"role": "assistant", "content": "Reply"}, + {"role": "user", "content": "Second"}, + ] + kwargs = build_converse_kwargs( + model="anthropic.claude-sonnet-4-6-20250514-v1:0", + messages=messages, + tools=tools, + ) + assert kwargs["system"][-1] == {"cachePoint": {"type": "default"}} + assert kwargs["toolConfig"]["tools"][-1] == {"cachePoint": {"type": "default"}} + # Second-to-last converse message (the assistant "Reply" turn) carries + # the checkpoint; the newest "Second" turn does not. + marked = kwargs["messages"][-2]["content"] + assert marked[-1] == {"cachePoint": {"type": "default"}} + assert kwargs["messages"][-1]["content"][-1] != {"cachePoint": {"type": "default"}} + + def test_no_cache_point_for_unsupported_model(self): + from agent.bedrock_adapter import build_converse_kwargs + messages = [ + {"role": "system", "content": "Be helpful."}, + {"role": "user", "content": "First"}, + {"role": "assistant", "content": "Reply"}, + {"role": "user", "content": "Second"}, + ] + kwargs = build_converse_kwargs(model="meta.llama3-70b-instruct-v1:0", messages=messages) + assert {"cachePoint": {"type": "default"}} not in kwargs["system"] + for m in kwargs["messages"]: + assert {"cachePoint": {"type": "default"}} not in m["content"] + # --------------------------------------------------------------------------- # Model discovery @@ -963,6 +1025,30 @@ class TestClientCache: class TestStreamConverseWithCallbacks: """Test real-time streaming with delta callbacks.""" + def test_cache_tokens_folded_into_prompt_tokens(self): + """The streaming path must fold cacheRead/WriteInputTokens into + prompt_tokens the same way the non-streaming path does (see + TestNormalizeConverseResponse.test_cache_tokens_folded_into_prompt_tokens).""" + from agent.bedrock_adapter import stream_converse_with_callbacks + events = {"stream": [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"contentBlockIndex": 0, "start": {}}}, + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hi"}}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + {"metadata": {"usage": { + "inputTokens": 50, + "outputTokens": 20, + "cacheReadInputTokens": 900, + "cacheWriteInputTokens": 300, + }}}, + ]} + result = stream_converse_with_callbacks(events) + assert result.usage.prompt_tokens == 50 + 900 + 300 + assert result.usage.total_tokens == 50 + 900 + 300 + 20 + assert result.usage.cache_read_input_tokens == 900 + assert result.usage.cache_creation_input_tokens == 300 + def test_text_deltas_fire_callback(self): from agent.bedrock_adapter import stream_converse_with_callbacks deltas = [] diff --git a/tests/agent/test_usage_pricing.py b/tests/agent/test_usage_pricing.py index a1f853272bb..3cfbbd231b1 100644 --- a/tests/agent/test_usage_pricing.py +++ b/tests/agent/test_usage_pricing.py @@ -25,6 +25,27 @@ def test_normalize_usage_anthropic_keeps_cache_buckets_separate(): assert normalized.prompt_tokens == 3400 +def test_normalize_usage_bedrock_converse_cache_point_round_trips(): + """End-to-end contract for the Converse cachePoint feature: the usage + shape bedrock_adapter.normalize_converse_response() produces (prompt_tokens + folded from inputTokens + cacheRead + cacheWrite, cache fields under their + Anthropic names) must normalize back to the original cache_read/write + split and the original Converse inputTokens value.""" + usage = SimpleNamespace( + prompt_tokens=50 + 900 + 300, + completion_tokens=20, + cache_read_input_tokens=900, + cache_creation_input_tokens=300, + ) + + normalized = normalize_usage(usage, provider="bedrock", api_mode="bedrock_converse") + + assert normalized.cache_read_tokens == 900 + assert normalized.cache_write_tokens == 300 + assert normalized.input_tokens == 50 + assert normalized.output_tokens == 20 + + def test_normalize_usage_openai_subtracts_cached_prompt_tokens(): usage = SimpleNamespace( prompt_tokens=3000,