mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(caching): stop counting an unrun server tool call mark, and follow #49192 and #35602 in their integration tests (#49468)
* fix(caching): stop counting an unrun server tool call mark, and follow #49192 and #35602 in their integration tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(caching): return the server tool call id instead of a TypeGuard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert the Gemini cache write carries the system prompt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang
GitHub
parent
8d202faec7
commit
3155aacda4
5 changed files with 72 additions and 44 deletions
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue