diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 06abb591717..8bb0e12905e 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3626,6 +3626,7 @@ class BedrockImageProcessor: def _convert_to_bedrock_tool_call_invoke( tool_calls: list, + model: Optional[str] = None, ) -> List[BedrockContentBlock]: """ OpenAI tool invokes: @@ -3701,7 +3702,13 @@ def _convert_to_bedrock_tool_call_invoke( # cache_control applies to the whole original # tool call; attach after the last split block. if tool.get("cache_control", None) is not None: - _parts_list.append(BedrockContentBlock(cachePoint=CachePointBlock(type="default"))) + _cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + {"cache_control": tool["cache_control"]}, + block_type="content_block", + model=model, + ) + if _cache_point_block is not None: + _parts_list.append(_cache_point_block) continue # Fallback: no objects extracted — use empty dict. arguments_dict = {} @@ -3712,8 +3719,13 @@ def _convert_to_bedrock_tool_call_invoke( # Check for cache_control and add a separate cachePoint block if tool.get("cache_control", None) is not None: - cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default")) - _parts_list.append(cache_point_block) + cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + {"cache_control": tool["cache_control"]}, + block_type="content_block", + model=model, + ) + if cache_point_block is not None: + _parts_list.append(cache_point_block) return _parts_list except Exception as e: raise Exception( @@ -4417,22 +4429,27 @@ class BedrockConverseMessagesProcessor: tool_content.append(tool_call_result) # Check if we need to add a separate cachePoint block - has_cache_control = False + tool_msg_cache_control = None # Check for message-level cache_control if current_message.get("cache_control", None) is not None: - has_cache_control = True + tool_msg_cache_control = current_message["cache_control"] # Check for content-level cache_control in list content elif isinstance(current_message.get("content"), list): for content_element in current_message["content"]: if isinstance(content_element, dict) and content_element.get("cache_control", None) is not None: - has_cache_control = True + tool_msg_cache_control = content_element["cache_control"] break # Add a separate cachePoint block if cache_control is present - if has_cache_control: - cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default")) - tool_content.append(cache_point_block) + if tool_msg_cache_control is not None: + cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + {"cache_control": tool_msg_cache_control}, + block_type="content_block", + model=model, + ) + if cache_point_block is not None: + tool_content.append(cache_point_block) msg_i += 1 # Deduplicate toolResult blocks with the same toolUseId @@ -4529,7 +4546,7 @@ class BedrockConverseMessagesProcessor: _tool_calls = assistant_message_block.get("tool_calls", []) if _tool_calls: - assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls)) + assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls, model=model)) msg_i += 1 @@ -4789,22 +4806,27 @@ def _bedrock_converse_messages_pt( tool_content.append(tool_call_result) # Check if we need to add a separate cachePoint block - has_cache_control = False + tool_msg_cache_control = None # Check for message-level cache_control if current_message.get("cache_control", None) is not None: - has_cache_control = True + tool_msg_cache_control = current_message["cache_control"] # Check for content-level cache_control in list content elif isinstance(current_message.get("content"), list): for content_element in current_message["content"]: if isinstance(content_element, dict) and content_element.get("cache_control", None) is not None: - has_cache_control = True + tool_msg_cache_control = content_element["cache_control"] break # Add a separate cachePoint block if cache_control is present - if has_cache_control: - cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default")) - tool_content.append(cache_point_block) + if tool_msg_cache_control is not None: + cache_point_block = litellm.AmazonConverseConfig()._get_cache_point_block( + {"cache_control": tool_msg_cache_control}, + block_type="content_block", + model=model, + ) + if cache_point_block is not None: + tool_content.append(cache_point_block) msg_i += 1 # Deduplicate toolResult blocks with the same toolUseId @@ -4902,7 +4924,7 @@ def _bedrock_converse_messages_pt( assistant_content.append(_cache_point_block) _tool_calls = assistant_message_block.get("tool_calls", []) if _tool_calls: - assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls)) + assistant_content.extend(_convert_to_bedrock_tool_call_invoke(_tool_calls, model=model)) msg_i += 1 diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index fe060e2e40d..9f2a4168dec 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -5671,3 +5671,99 @@ async def test_grounding_source_and_query_rendered_as_text(): user_content = result[0]["content"] assert {"text": "Tokyo is the capital of Japan."} in user_content assert {"text": "What is the capital of Japan?"} in user_content + + +def _agentic_messages_with_ttl(ttl_target: str): + """A tool-loop conversation with `ttl: 1h` cache_control at `ttl_target`: + 'user', 'tool_call' (per-tool-call, on the assistant's tool call), or + 'tool' (message-level, on the tool result - where + `cache_control_injection_points` with `index: -1` lands mid-loop). + + Message-level cache_control on a content-less assistant message emits no + cachePoint at all today (a separate gap, orthogonal to ttl); per-tool-call + placement covers that message, so it's excluded from the params below.""" + user: dict = {"role": "user", "content": "optimize this kernel " * 60} + assistant: dict = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "evaluate", "arguments": "{}"}, + } + ], + } + tool: dict = {"role": "tool", "tool_call_id": "call_1", "content": "score: 42"} + ttl_cc = {"type": "ephemeral", "ttl": "1h"} + if ttl_target == "user": + user["cache_control"] = ttl_cc + elif ttl_target == "tool_call": + assistant["tool_calls"][0]["cache_control"] = ttl_cc + elif ttl_target == "tool": + tool["cache_control"] = ttl_cc + return [user, assistant, tool] + + +def _collect_cache_points(result): + return [ + block["cachePoint"] + for message in result + for block in message.get("content") or [] + if "cachePoint" in block + ] + + +@pytest.mark.parametrize("ttl_target", ["user", "tool_call", "tool"]) +@pytest.mark.asyncio +async def test_message_level_cache_control_honors_ttl_for_supported_model( + ttl_target, +): + """Message- and tool-call-level cache_control must carry `ttl` onto the + emitted cachePoint for models that support extended caching, mirroring the + system-message path. Regression test for the gap left by the system-only + fix: the message paths called `_get_cache_point_block` without `model` (or + hardcoded `{"type": "default"}`), silently downgrading 1h to 5m.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + BedrockConverseMessagesProcessor, + _bedrock_converse_messages_pt, + ) + + messages = _agentic_messages_with_ttl(ttl_target) + + result = _bedrock_converse_messages_pt( + messages=messages, + model="global.anthropic.claude-opus-4-7", + llm_provider="bedrock_converse", + ) + async_result = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model="global.anthropic.claude-opus-4-7", + llm_provider="bedrock_converse", + ) + ) + assert result == async_result + + cache_points = _collect_cache_points(result) + assert len(cache_points) == 1 + assert cache_points[0].get("ttl") == "1h" + + +@pytest.mark.parametrize("ttl_target", ["user", "tool_call", "tool"]) +def test_message_level_cache_control_drops_ttl_for_unsupported_model(ttl_target): + """Models outside the extended-caching allow-list must keep emitting the + plain `{"type": "default"}` cachePoint (Bedrock rejects `ttl` for them).""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _bedrock_converse_messages_pt, + ) + + result = _bedrock_converse_messages_pt( + messages=_agentic_messages_with_ttl(ttl_target), + model="anthropic.claude-3-5-sonnet-20240620-v1:0", + llm_provider="bedrock_converse", + ) + + cache_points = _collect_cache_points(result) + assert len(cache_points) == 1 + assert "ttl" not in cache_points[0]