diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index faedf8ae1a3..a0f7da074e0 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -10,6 +10,7 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and """ import copy +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_logger @@ -34,6 +35,9 @@ else: # breakpoints: "A maximum of 4 blocks with cache_control may be provided." MAX_CACHE_CONTROL_BLOCKS = 4 +# Providers whose transform turns a tool-level cache_control into a provider cache breakpoint +TOOL_CACHE_CONTROL_PROVIDERS = frozenset(("bedrock", "bedrock_converse")) + class AnthropicCacheControlHook(CustomPromptManagement): def get_chat_completion_prompt( @@ -298,6 +302,51 @@ class AnthropicCacheControlHook(CustomPromptManagement): return processed_messages, processed_system, remaining_points + @staticmethod + def apply_tool_config_injection_points( + tools: Sequence[Mapping[str, Any]] | None, + non_default_params: Mapping[str, Any], + custom_llm_provider: str | None, + ) -> tuple[Sequence[Mapping[str, Any]] | None, Mapping[str, Any]]: + """Move a Bedrock ``tool_config`` injection point onto the tools themselves. + + The Bedrock transform turns a tool-level ``cache_control`` into a + ``cachePoint`` tool block, so carrying the breakpoint on the tool + produces the same request while making it visible to logging callbacks, + which log ``optional_params["tools"]`` (see issue #34758). Returns the + tools and params to use; the tool_config point is dropped from the + params it stamps so the transform doesn't append a second cachePoint. + + Stands down when the trailing tool is a pre-formatted Bedrock tool block + (e.g. ``systemTool``), which the transform passes through untouched and + would therefore drop the stamp, and when the client already marked that + tool; in both cases the transform's trailing cachePoint still applies. + """ + points: Sequence[CacheControlInjectionPoint] = non_default_params.get("cache_control_injection_points") or () + if not tools or custom_llm_provider not in TOOL_CACHE_CONTROL_PROVIDERS: + return tools, non_default_params + + tool_config_points = tuple(point for point in points if point.get("location") == "tool_config") + if not tool_config_points: + return tools, non_default_params + + last_tool = tools[-1] + if not isinstance(last_tool, dict) or not ("function" in last_tool or "input_schema" in last_tool): + return tools, non_default_params + + remaining_points = tuple(point for point in points if point.get("location") != "tool_config") + updated_params = ( + {**non_default_params, "cache_control_injection_points": remaining_points} + if remaining_points + else {key: value for key, value in non_default_params.items() if key != "cache_control_injection_points"} + ) + + if last_tool.get("cache_control") is not None: + return tools, updated_params + + control = {**(tool_config_points[0].get("control") or {}), "type": "ephemeral"} + return [*tools[:-1], {**last_tool, "cache_control": control}], updated_params + @staticmethod def _default_control() -> ChatCompletionCachedContent: """Build the cache_control block for auto-injected breakpoints. diff --git a/litellm/main.py b/litellm/main.py index dc3ec469a1b..9a0295b5353 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5161,6 +5161,17 @@ def completion( # type: ignore if not _should_allow_input_examples(custom_llm_provider=custom_llm_provider, model=model): tools = _drop_input_examples_from_tools(tools=tools) + ( + tools_with_cache_control, + params_with_cache_control, + ) = AnthropicCacheControlHook.apply_tool_config_injection_points( + tools=tools, + non_default_params=non_default_params, + custom_llm_provider=custom_llm_provider, + ) + tools = None if tools_with_cache_control is None else [*tools_with_cache_control] + non_default_params = {**params_with_cache_control} + if provider_specific_header is not None: headers.update( ProviderSpecificHeaderUtils.get_provider_specific_headers( diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index d94f0d5f47e..76f2a8e5db0 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1,3 +1,4 @@ +import asyncio import copy import datetime import json @@ -15,6 +16,7 @@ import pytest sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system-path import litellm from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import StandardCallbackDynamicParams @@ -1346,6 +1348,150 @@ async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(): ) +@pytest.mark.asyncio +async def test_tool_config_injection_point_visible_to_logging_callbacks(): + """Regression test for #34758. + + A tool_config injection point used to be applied only inside the Bedrock + transform, so logging callbacks (which log ``optional_params["tools"]``) + never saw the tool-level breakpoint, unlike message-level cache_control. + The stamped tool must reach the callbacks, and the Bedrock payload must + still carry exactly one tool cachePoint (no double injection). + """ + logged_kwargs: dict = {} + + class CaptureLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + logged_kwargs.update(kwargs) + + with patch.dict( + os.environ, + { + "AWS_ACCESS_KEY_ID": "fake_access_key_id", + "AWS_SECRET_ACCESS_KEY": "fake_secret_access_key", + "AWS_REGION_NAME": "us-east-1", + }, + ): + litellm.callbacks = [CaptureLogger()] + + mock_response = MagicMock() + mock_response.json.return_value = { + "output": {"message": {"role": "assistant", "content": "ok"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104}, + } + mock_response.status_code = 200 + + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=mock_response) as mock_post: + await litellm.acompletion( + model="bedrock/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "What is the weather?"}], + max_tokens=32, + tools=[ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ], + cache_control_injection_points=[ + {"location": "tool_config", "control": {"type": "ephemeral", "ttl": "1h"}}, + ], + client=client, + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + + request_tools = request_body["toolConfig"]["tools"] + assert [tool for tool in request_tools if "cachePoint" in tool] == [ + {"cachePoint": {"type": "default", "ttl": "1h"}} + ] + + for _ in range(100): + if logged_kwargs: + break + await asyncio.sleep(0.1) + + logged_tools = logged_kwargs["optional_params"]["tools"] + assert logged_tools[-1]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +def test_apply_tool_config_injection_points_leaves_non_bedrock_providers_alone(): + """Only Bedrock turns a tool-level cache_control into a provider cache breakpoint. + + Stamping tools for other providers would forward an unsupported field + upstream, so the point stays untouched for them. + """ + tools = [{"type": "function", "function": {"name": "get_weather"}}] + non_default_params = {"cache_control_injection_points": [{"location": "tool_config"}]} + + updated_tools, updated_params = AnthropicCacheControlHook.apply_tool_config_injection_points( + tools=tools, + non_default_params=non_default_params, + custom_llm_provider="openai", + ) + + assert updated_tools == tools + assert updated_params == non_default_params + + +def test_apply_tool_config_injection_points_keeps_message_points(): + tools = [{"type": "function", "function": {"name": "get_weather"}}] + message_point = {"location": "message", "role": "system"} + + updated_tools, updated_params = AnthropicCacheControlHook.apply_tool_config_injection_points( + tools=tools, + non_default_params={"cache_control_injection_points": [message_point, {"location": "tool_config"}]}, + custom_llm_provider="bedrock", + ) + + assert updated_tools[-1]["cache_control"] == {"type": "ephemeral"} + assert list(updated_params["cache_control_injection_points"]) == [message_point] + + +def test_apply_tool_config_injection_points_defers_to_client_marked_tool(): + """A client-marked tool already yields a cachePoint; don't add a second one.""" + tools = [ + {"type": "function", "function": {"name": "get_weather"}, "cache_control": {"type": "ephemeral"}}, + ] + + updated_tools, updated_params = AnthropicCacheControlHook.apply_tool_config_injection_points( + tools=tools, + non_default_params={"cache_control_injection_points": [{"location": "tool_config"}]}, + custom_llm_provider="bedrock", + ) + + assert updated_tools == tools + assert "cache_control_injection_points" not in updated_params + + +def test_apply_tool_config_injection_points_defers_to_transform_for_bedrock_tool_blocks(): + """Pre-formatted Bedrock tool blocks pass the transform untouched. + + A stamp on them would be dropped, so the point must survive for the + transform to append its trailing cachePoint. + """ + tools = [{"systemTool": {"name": "nova_grounding"}}] + non_default_params = {"cache_control_injection_points": [{"location": "tool_config"}]} + + updated_tools, updated_params = AnthropicCacheControlHook.apply_tool_config_injection_points( + tools=tools, + non_default_params=non_default_params, + custom_llm_provider="bedrock", + ) + + assert updated_tools == tools + assert updated_params == non_default_params + + class TestApplyToAnthropicMessagesRequest: """Tests for apply_to_anthropic_messages_request (v1/messages cache control)."""