mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-31 19:16:29 +00:00
feat(bedrock): add Converse API prompt caching (cachePoint)
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.
This commit is contained in:
joaomarcos
Teknium
parent
12096b1e3d
commit
4dccfcd9b7
3 changed files with 168 additions and 11 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue