diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 7df13541828..fbefce48293 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -313,6 +313,14 @@ def _builds_anthropic_messages(model: str, custom_llm_provider: str | None) -> b return isinstance(chat_config, AnthropicConfig) and not isinstance(chat_config, OpenAIGPTConfig) +def _tool_call_breakpoints_become_marks( + messages: Sequence[object], model: str, custom_llm_provider: str | None +) -> bool: + if not any(_carries_mark(tool_call, "prompt_cache_breakpoint") for tool_call in _tool_calls(messages)): + return True + return _builds_anthropic_messages(model, custom_llm_provider) or _served_by_bedrock(model, custom_llm_provider) + + def _custom_llm_provider(non_default_params: Mapping[str, object], point: Mapping[str, object]) -> str | None: provider: Final = non_default_params.get("custom_llm_provider") or point.get(CUSTOM_LLM_PROVIDER_STAMP) return provider if isinstance(provider, str) else None @@ -475,8 +483,11 @@ class AnthropicCacheControlHook(CustomPromptManagement): answered_tool_ids: Final = _tool_call_ids_sent_as_tool_calls_when_a_server_tool_call_is_marked( processed_messages, model, custom_llm_provider ) + census_dialect: Final = openai_dialect or not _tool_call_breakpoints_become_marks( + processed_messages, model, custom_llm_provider + ) breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints( - processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=openai_dialect + processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=census_dialect ) processed_messages = self._apply_message_injections( points=applied_message_points, @@ -484,6 +495,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks, openai_dialect=openai_dialect, answered_tool_ids=answered_tool_ids, + census_dialect=census_dialect, ) if ( openai_dialect @@ -504,7 +516,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): *AnthropicCacheControlHook._points_with_a_slot_left( remaining_points, AnthropicCacheControlHook.count_request_cache_breakpoints( - processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=openai_dialect + processed_messages, answered_tool_ids=answered_tool_ids, openai_dialect=census_dialect ) + external_breakpoints, openai_dialect, @@ -634,6 +646,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): max_blocks: int, openai_dialect: bool = False, answered_tool_ids: frozenset[str] | None = None, + census_dialect: bool | None = None, ) -> list[AllMessageValues]: """Apply message-level cache control injection points in order. @@ -649,8 +662,9 @@ class AnthropicCacheControlHook(CustomPromptManagement): if answered_tool_ids is None else answered_tool_ids ) + counted_dialect: Final = openai_dialect if census_dialect is None else census_dialect used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints( - messages, answered_tool_ids=client_answers, openai_dialect=openai_dialect + messages, answered_tool_ids=client_answers, openai_dialect=counted_dialect ) limit_reached = False @@ -669,7 +683,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): break if AnthropicCacheControlHook._message_has_cache_control( - messages[target_index], client_answers, openai_dialect + messages[target_index], client_answers, counted_dialect ): # Client already marked this message; don't overwrite it. continue @@ -678,7 +692,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): messages[target_index], control, openai_dialect ) if AnthropicCacheControlHook._message_has_cache_control( - messages[target_index], client_answers, openai_dialect + messages[target_index], client_answers, counted_dialect ): used_blocks += 1 diff --git a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py index 0b4081b8348..8baa1fd57a2 100644 --- a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py +++ b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py @@ -152,9 +152,13 @@ def _namespace_members(tools: Sequence[object]) -> Iterator[tuple[_Path, _Fields ) +def _function_mark(tool: _Fields) -> _Fields | None: + return _fields((_fields(tool.get("function")) or {}).get("cache_control")) + + def _member_mark(member: _Fields) -> object: own: Final = member.get("cache_control") - return own if own is not None else (_fields(member.get("function")) or {}).get("cache_control") + return own if own is not None else _function_mark(member) def _transform_sends_member(member: _Fields) -> bool: @@ -530,12 +534,13 @@ def _on_wire( ) -> bool: """Claude sends a message-level mark on string content, or on list content only for a tool result, and none on a developer message, which reaches it as a system message without one; a tool's mark only when the route sends - it, and one under ``function`` only off Converse and when the tool carries none of its own; no tool_search call - nobody answered; and no mark on tool history the route sends as text.""" + it, and one under ``function`` only off Converse, as an object, and when the tool carries none of its own; no + tool_search call nobody answered; and no mark on tool history the route sends as text.""" if path[0] == "tools": tool: Final = _fields(tools[int(path[1])]) or {} return tool_mark_reaches_claude(tool, route) and ( - len(path) == 2 or (route != "converse" and tool.get("cache_control") is None) + len(path) == 2 + or (route != "converse" and tool.get("cache_control") is None and _function_mark(tool) is not None) ) message: Final = _fields(messages[int(path[1])]) or {} if tool_history_as_text and _is_tool_history(message): diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 8d83c281bbd..d984edddcec 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -44,7 +44,7 @@ _MARK_KEYS: Final = frozenset({"cache_control", "prompt_cache_breakpoint"}) _ENVELOPE_SKIPPED_KEYS: Final = _MARK_KEYS | {"content"} _PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) _TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...]) -_OBJECTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None) @@ -79,8 +79,12 @@ def _carries_cache_mark(block: Mapping[str, object]) -> bool: def _carries_marked_tool_call(message: Mapping[str, object]) -> bool: - tool_calls: Final = _OBJECTS_ADAPTER.validate_python(message.get("tool_calls") or ()) - return any(_carries_cache_mark(call) for call in tool_calls) + tool_calls: Final = _JSON_ADAPTER.validate_python( + to_jsonable_python(message.get("tool_calls"), serialize_unknown=True, bytes_mode="base64") + ) + return isinstance(tool_calls, list) and any( + isinstance(call, dict) and _carries_cache_mark(call) for call in tool_calls + ) def _unmarked(block: JsonValue) -> JsonValue: diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 018fa808f70..e708b0216e7 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -19,6 +19,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ) from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter MODEL_GROUP_ALIAS = "my-claude-group" @@ -860,6 +861,48 @@ async def test_pin_survives_the_marked_message_coming_back_as_string_content(): assert filtered == [deployments[0]] +def _lookup_call() -> ChatCompletionMessageToolCall: + return ChatCompletionMessageToolCall( + id="call_1", type="function", function=Function(name="lookup_clause", arguments="{}") + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "assistant_turn", + [ + Message(content=None, tool_calls=[_lookup_call()]), + {"role": "assistant", "content": None, "tool_calls": [_lookup_call()]}, + {"role": "assistant", "content": None, "tool_calls": "call_1"}, + {"role": "assistant", "content": None, "tool_calls": {"id": "call_1"}}, + {"role": "assistant", "content": None, "tool_calls": [1]}, + {"role": "assistant", "content": None, "tool_calls": ["call_1"]}, + {"role": "assistant", "content": None, "tool_calls": [None]}, + ], + ids=["response_message", "response_tool_calls", "string", "object", "number_list", "string_list", "null_list"], +) +async def test_pin_holds_whatever_shape_the_assistant_tool_calls_come_back_in(assistant_turn: object): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + turn_one = _turn( + {"role": "system", "content": [_marked(LONG_PROMPT)]}, + {"role": "user", "content": "Find the termination clause."}, + ) + turn_two = cast( + list[AllMessageValues], + [*turn_one, assistant_turn, {"role": "tool", "tool_call_id": "call_1", "content": "Section 9."}], + ) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=turn_one, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[1]] + + @pytest.mark.asyncio async def test_lookback_stops_where_the_provider_cache_stops(): """ diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 5d07abd0111..3e54111a8c8 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -1148,6 +1148,34 @@ def test_a_breakpoint_drop_params_drops_leaves_key_auto_caching_on(respx_mock: r {"type": "tool_search_tool_regex_20251119", "name": "tool_search", "cache_control": _FIVE_MINUTES}, id="converse-tool-search-regex", ), + pytest.param( + "anthropic/claude-opus-5-5", + r"https://api\.anthropic\.com/v1/messages", + _CLAUDE_REPLY, + {"api_key": "fake"}, + { + "type": "function", + "function": {"name": "spawn", "parameters": {"type": "object"}, "cache_control": "ephemeral"}, + }, + id="anthropic-string-function-mark", + ), + pytest.param( + "bedrock/invoke/us.anthropic.claude-opus-5-5", + _INVOKE_ROUTE, + _CLAUDE_REPLY, + _FAKE_AWS, + { + "type": "namespace", + "name": "agents", + "tools": [ + { + "type": "function", + "function": {"name": "spawn", "parameters": {"type": "object"}, "cache_control": "ephemeral"}, + } + ], + }, + id="invoke-namespace-member-string-function-mark", + ), ], ) def test_a_five_minute_mark_on_a_tool_the_route_sends_unmarked_leaves_the_system_breakpoint_free_to_take_the_hour( @@ -1283,6 +1311,60 @@ def test_a_marked_past_server_tool_call_leaves_the_operator_system_point_in_plac assert "cache_control" in json.dumps(json.loads(route.calls.last.request.content)["messages"][0]) +@pytest.mark.parametrize( + ("model", "api_base", "url"), + [ + ("openrouter/anthropic/claude-opus-5-5", None, "https://openrouter.ai/api/v1/chat/completions"), + ( + "databricks/databricks-claude-opus-4-5", + "https://adb-cache.cloud.databricks.com/serving-endpoints", + "https://adb-cache.cloud.databricks.com/serving-endpoints/chat/completions", + ), + ], + ids=["openrouter", "databricks"], +) +def test_a_breakpoint_on_a_tool_call_the_route_never_translates_leaves_the_operator_system_point_in_place( + respx_mock: respx.MockRouter, + openai_api_response: Mapping[str, object], + model: str, + api_base: str | None, + url: str, +): + route: Final = respx_mock.post(url).respond(json=openai_api_response) + + litellm.completion( + model=model, + messages=[ + {"role": "system", "content": "sys"}, + { + "role": "user", + "content": [ + {"type": "text", "text": str(index), "prompt_cache_breakpoint": _BREAKPOINT} for index in range(3) + ], + }, + { + "role": "assistant", + "content": "Looking it up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + "prompt_cache_breakpoint": _BREAKPOINT, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Sunny"}, + {"role": "user", "content": "Thanks"}, + ], + cache_control_injection_points=[{"location": "message", "role": "system"}], + api_key="fake", + api_base=api_base, + ) + + assert "cache_control" in json.dumps(json.loads(route.calls.last.request.content)["messages"][0]) + + def test_a_tuple_additional_drop_params_places_the_same_breakpoints_as_a_list(respx_mock: respx.MockRouter): route: Final = respx_mock.post(url__regex=_CONVERSE_ROUTE).respond(json=_CONVERSE_REPLY)