diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 39368dee47c..fa519740e88 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -13,7 +13,7 @@ import copy import itertools import os import re -from collections.abc import Callable, Container, Iterable, Mapping, Sequence +from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast from urllib.parse import urlparse @@ -445,22 +445,41 @@ def _as_object_list(value: object | None) -> list[object] | None: return _validated_object_list(value) -def _tool_call_carries_cache_breakpoint(tool_call: object, message: object, answered_tool_ids: Container[str]) -> bool: +def _server_tool_call_id(tool_call: object) -> str | None: + tool_call_id: Final = _attribute_or_key(tool_call, "id") + if isinstance(tool_call_id, str) and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX): + return tool_call_id + return None + + +def _tool_call_carries_cache_breakpoint(tool_call: object, client_tool_call_ids: Container[str]) -> bool: if _attribute_or_key(tool_call, "cache_control") is None: return False + server_tool_call_id: Final = _server_tool_call_id(tool_call) + return server_tool_call_id is None or server_tool_call_id in client_tool_call_ids - return not tool_call_is_rebuilt_as_server_tool_use( - _attribute_or_key(tool_call, "id"), _attribute_or_key(message, "provider_specific_fields"), answered_tool_ids + +def _server_tool_call_ids_without_a_kept_result(message: object) -> Iterator[str]: + provider_specific_fields: Final = _attribute_or_key(message, "provider_specific_fields") + tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls")) or () + return ( + tool_call_id + for tool_call in tool_calls + if (tool_call_id := _server_tool_call_id(tool_call)) is not None + and not tool_call_is_rebuilt_as_server_tool_use(tool_call_id, provider_specific_fields) ) -def _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked( +def _client_tool_call_ids_when_a_server_tool_call_is_marked( messages: Sequence[object], model: str | None = None, custom_llm_provider: str | None = None ) -> frozenset[str]: if not _marks_a_server_tool_call(messages): return frozenset() if model is not None and not _builds_anthropic_messages(model, custom_llm_provider): - return tool_call_ids_answered_in(messages) + server_calls_sent_as_client_calls: Final = itertools.chain.from_iterable( + map(_server_tool_call_ids_without_a_kept_result, messages) + ) + return tool_call_ids_answered_in(messages) | frozenset(server_calls_sent_as_client_calls) return tool_call_ids_answered_by_client(messages, _keeps_mid_conversation_system(model, custom_llm_provider)) @@ -517,12 +536,7 @@ def _model_info_flags_mid_conversation_system(model: str, custom_llm_provider: s def _is_marked_server_tool_call(tool_call: object) -> bool: - tool_call_id: Final = _attribute_or_key(tool_call, "id") - return ( - isinstance(tool_call_id, str) - and tool_call_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX) - and _attribute_or_key(tool_call, "cache_control") is not None - ) + return _server_tool_call_id(tool_call) is not None and _attribute_or_key(tool_call, "cache_control") is not None def _tool_carries_cache_breakpoint(tool: object) -> bool: @@ -668,7 +682,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages( remaining_points, external_breakpoints, openai_dialect ) - answered_tool_ids: Final = _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked( + answered_tool_ids: Final = _client_tool_call_ids_when_a_server_tool_call_is_marked( processed_messages, model, custom_llm_provider ) breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints( @@ -757,7 +771,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) message_list: Final = tuple(messages) client_answers: Final = ( - _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(message_list) + _client_tool_call_ids_when_a_server_tool_call_is_marked(message_list) if answered_tool_ids is None else answered_tool_ids ) @@ -849,7 +863,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): so earlier points win when slots are scarce. """ client_answers: Final = ( - _tool_call_ids_answered_by_client_when_a_server_tool_call_is_marked(messages) + _client_tool_call_ids_when_a_server_tool_call_is_marked(messages) if answered_tool_ids is None else answered_tool_ids ) @@ -939,11 +953,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): content_count: Final = sum(1 for block in content if _carries_cache_breakpoint(block)) if content else 0 tool_calls: Final = _as_object_list(_attribute_or_key(message, "tool_calls")) tool_call_count: Final = ( - sum( - 1 - for tool_call in tool_calls - if _tool_call_carries_cache_breakpoint(tool_call, message, answered_tool_ids) - ) + sum(1 for tool_call in tool_calls if _tool_call_carries_cache_breakpoint(tool_call, answered_tool_ids)) if tool_calls else 0 ) diff --git a/tests/integration/_support/anthropic_refusal.py b/tests/integration/_support/anthropic_refusal.py index 6d5fd236074..fb6c5e71a82 100644 --- a/tests/integration/_support/anthropic_refusal.py +++ b/tests/integration/_support/anthropic_refusal.py @@ -29,6 +29,7 @@ DOCS_STREAMING_REFUSAL_URL: Final = ( ) # Copied verbatim from DOCS_REFUSAL_URL and DOCS_STREAMING_REFUSAL_URL on 2026-10-12. +DOCS_REFUSAL_EXPLANATION: Final = "This request was declined because it could enable cyber harm." DOCS_REFUSAL: Final[Event] = { "id": "msg_01XFUDYJgAACzvnptvVoYEL", "type": "message", @@ -39,7 +40,7 @@ DOCS_REFUSAL: Final[Event] = { "stop_details": { "type": "refusal", "category": "cyber", - "explanation": "This request was declined because it could enable cyber harm.", + "explanation": DOCS_REFUSAL_EXPLANATION, }, "usage": {"input_tokens": 412, "output_tokens": 0}, } @@ -50,7 +51,7 @@ DOCS_PARTIAL_REFUSAL: Final[Event] = { "stop_details": { "type": "refusal", "category": "cyber", - "explanation": "This request was declined because it could enable cyber harm.", + "explanation": DOCS_REFUSAL_EXPLANATION, }, } diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py index cd42e816ee3..c9a85545d8b 100644 --- a/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_wire.py @@ -18,6 +18,7 @@ from integration._support.anthropic_refusal import ( DECLARED_PARTIAL_JSON, DOCS_PARTIAL_TEXT, DOCS_REFUSAL, + DOCS_REFUSAL_EXPLANATION, DOCS_STOP_DETAILS, ENVELOPED_PARTIAL_JSON, FIXTURES, @@ -651,8 +652,7 @@ def test_family_messages_stream_refusal_keeps_frames_and_logs_the_stop_details(r model: Final = _deployment(scenario, provider, rig) key: Final = scenario.key() events: Final = _messages_stream(rig, key, model, refusal_prompt("refusal", _marker())) - native: Final = {"stop_reason": "refusal"} if provider == "snowflake" else REFUSAL_FIELDS - assert _native_deltas(events) == (native,), events + assert _native_deltas(events) == (REFUSAL_FIELDS,), events assert _relayed_in(events) == (), events assert _paths(rig.wire) == (_STREAM_TARGETS[provider][1],) logged: Final = _logged_response(rig, key) @@ -783,10 +783,10 @@ def test_responses_cache_hit_logs_the_same_stop_details_as_the_miss(rig: Rig) -> served_from_cache: Final = tuple(event.get("cache_hit") is True for event in logged) assert (len(_paths(rig.wire)), served_from_cache.count(True)) == (1, 1), (served_from_cache, logged) assert relayed(first) == ((_RESPONSE_FIELDS, REFUSAL_FIELDS),), first - assert tuple(relayed(event["response"]) for event in logged) == (((_CHAT_FIELDS, REFUSAL_FIELDS),),) * 2, ( - served_from_cache, - logged, - ) + assert tuple(relayed(event["response"]) for event in logged) == ( + ((_CHAT_FIELDS, REFUSAL_FIELDS),), + relayed(first), + ), (served_from_cache, logged) def test_passthrough_is_not_served_from_the_response_cache(rig: Rig) -> None: @@ -1310,6 +1310,7 @@ def test_chat_refusal_body_serializes_the_stop_details_beside_the_existing_messa "message": { "content": content, "role": "assistant", + "refusal": DOCS_REFUSAL_EXPLANATION, "provider_specific_fields": {**_ANTHROPIC_MESSAGE_FIELDS, **REFUSAL_FIELDS}, }, } diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py index a18224a9f36..9fc13c77bcc 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_wire.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_wire.py @@ -354,7 +354,7 @@ _WEB_RESULTS: Final[dict[str, JsonValue]] = { ), pytest.param( {}, - [ASK_LABEL, tool_use_label(CITIES[0]), tool_use_label(CITIES[1]), "assistant:tool_use:srvtoolu_web"], + [SYSTEM_LABEL, ASK_LABEL, tool_use_label(CITIES[0]), tool_use_label(CITIES[1])], id="server-tool-without-result", ), ), @@ -413,16 +413,14 @@ _GEMINI_CACHE_WRITE: Final = [ @pytest.mark.parametrize( - ("calls", "expected"), + "calls", ( - pytest.param( - marked_calls(), [("POST", "/models/gemini-3.8-flash:generateContent")], id="tool-calls-fill-the-cap" - ), - pytest.param([tool_call(city) for city in CITIES], _GEMINI_CACHE_WRITE, id="unmarked-tool-calls"), + pytest.param(marked_calls(), id="tool-calls-fill-the-cap"), + pytest.param([tool_call(city) for city in CITIES], id="unmarked-tool-calls"), ), ) -def test_gemini_context_cache_follows_the_cap_census( - gateway: Gateway, calls: list[JsonValue], expected: list[tuple[str, str]] +def test_gemini_context_cache_takes_the_system_message_with_the_marked_ask( + gateway: Gateway, calls: list[JsonValue] ) -> None: marker: Final = new_marker() long_system: Final[JsonValue] = {"role": "system", "content": "lorem " * (2 * get_prompt_cache_min_tokens(_GEMINI))} @@ -436,8 +434,18 @@ def test_gemini_context_cache_follows_the_cap_census( ) status, _, text = post_chat(gateway, chat_body(model, messages)) assert status == 200, text - targets: Final = [(request.method, request.target.split("?")[0]) for request in wire.drain()] - assert targets == expected + received: Final = wire.drain() + targets: Final = [(request.method, request.target.split("?")[0]) for request in received] + cache_writes: Final = [ + request.body for request in received if request.method == "POST" and "cachedContents" in request.target + ] + generate: Final = _JSON_OBJECT.validate_json(received[-1].body) + assert ( + targets, + [b"lorem lorem" in body for body in cache_writes], + generate.get("cachedContent"), + "system_instruction" in generate, + ) == (_GEMINI_CACHE_WRITE, [True], "cachedContents/census", False), (cache_writes, generate) def _openai_marks(request: Request) -> list[str]: @@ -747,11 +755,11 @@ def test_request_level_points_skip_injection_when_client_marks_fill_the_cap(gate _OTHER_CALL_RESULT: Final[dict[str, JsonValue]] = { "web_search_results": [{"type": "web_search_tool_result", "tool_use_id": "srvtoolu_other", "content": []}] } -_SERVER_MARK_COUNTED: Final = [ +_UNRUN_SERVER_MARK_UNCOUNTED: Final = [ + SYSTEM_LABEL, ASK_LABEL, tool_use_label(CITIES[0]), tool_use_label(CITIES[1]), - "assistant:tool_use:srvtoolu_web", ] @@ -767,7 +775,7 @@ _SERVER_MARK_COUNTED: Final = [ pytest.param(_OTHER_CALL_RESULT, id="result-for-another-call"), ), ) -def test_malformed_provider_specific_fields_count_the_server_tool_call_mark( +def test_malformed_provider_specific_fields_leave_the_unrun_server_tool_call_mark_uncounted( gateway: Gateway, fields: JsonValue ) -> None: marker: Final = new_marker() @@ -779,7 +787,7 @@ def test_malformed_provider_specific_fields_count_the_server_tool_call_mark( ) assert status == 200, text labels: Final = anthropic_labels(_only_request(wire)) - assert labels == _SERVER_MARK_COUNTED + assert labels == _UNRUN_SERVER_MARK_UNCOUNTED @pytest.mark.parametrize( diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index fa03879c05a..05a63030c69 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -1116,7 +1116,7 @@ def test_cache_control_hook_counts_tool_call_cache_controls(): assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 3 -def test_cache_control_hook_counts_tool_call_marks_except_answered_server_tool_calls(): +def test_cache_control_hook_counts_tool_call_marks_except_server_tool_calls_the_client_did_not_answer(): message: Final[AllMessageValues] = { "role": "assistant", "content": None, @@ -1170,7 +1170,7 @@ def test_cache_control_hook_counts_tool_call_marks_except_answered_server_tool_c }, } - assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 2 + assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 1 _SEARCH_ANSWER: Final = {"role": "tool", "tool_call_id": "srvtoolu_01Search", "content": "found get_weather"} @@ -1182,6 +1182,12 @@ _SEARCH_CALL: Final = { "function": {"name": "tool_search_tool_bm25", "arguments": "{}"}, "cache_control": {"type": "ephemeral"}, } +_UNRUN_WEB_SEARCH_CALL: Final = { + "id": "srvtoolu_01Web", + "type": "function", + "function": {"name": "web_search", "arguments": "{}"}, + "cache_control": {"type": "ephemeral"}, +} _WEATHER_CALL: Final = { "id": "toolu_01Weather", "type": "function", @@ -1238,6 +1244,8 @@ def _replayed_search_call(*tool_calls: dict[str, object]) -> AllMessageValues: (_SEARCH_CALL,), (_SEARCH_ANSWER, _GO_ON), False, 0, id="client_answer_replaced_by_the_kept_result" ), pytest.param((_SEARCH_CALL,), (_GO_ON,), True, 0, id="placeholder_replaced_by_the_kept_result"), + pytest.param((_UNRUN_WEB_SEARCH_CALL,), (), False, 0, id="trailing_unrun_search_replayed_without_its_mark"), + pytest.param((_UNRUN_WEB_SEARCH_CALL,), (_GO_ON,), False, 0, id="unrun_search_dropped_before_a_new_message"), pytest.param( (_SEARCH_CALL, _WEATHER_CALL), (_SEARCH_ANSWER,),