mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(router): read tool call marks without walking their values
The router's prompt caching check now reads each tool call's own cache_control and prompt_cache_breakpoint fields instead of serializing the whole tool call, so a tool call that carries a deeply nested value routes as before instead of answering 500. The cache control hook no longer asks the provider resolver whether a GitHub Copilot or ChatGPT model is served by Bedrock, so a breakpoint on a tool call never starts their device login. Tests cover tool calls in response objects and tuples, a function form namespace member mark, the Copilot login, and the two 30m e2e cases now set drop_params on their own deployment
This commit is contained in:
parent
cd17ced013
commit
3f1e289711
6 changed files with 131 additions and 18 deletions
|
|
@ -283,9 +283,11 @@ def _marks_a_server_tool_call(messages: Sequence[object]) -> bool:
|
|||
|
||||
def _served_by_bedrock(model: str, custom_llm_provider: str | None) -> bool:
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
if declared_authenticating_provider(model, custom_llm_provider) is not None:
|
||||
return False
|
||||
try:
|
||||
_, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except (BadRequestError, ValueError):
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from dataclasses import dataclass
|
|||
from itertools import accumulate, chain
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
from pydantic_core import to_jsonable_python
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
|
|
@ -44,7 +44,6 @@ _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, ...])
|
||||
_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None)
|
||||
|
||||
|
||||
|
|
@ -78,12 +77,18 @@ def _carries_cache_mark(block: Mapping[str, object]) -> bool:
|
|||
return _mark_pin_ttl_seconds(block) is not None
|
||||
|
||||
|
||||
def _tool_call_fields(call: object) -> Mapping[str, object]:
|
||||
if isinstance(call, BaseModel):
|
||||
return dict(call)
|
||||
if isinstance(call, Mapping):
|
||||
return call
|
||||
return {}
|
||||
|
||||
|
||||
def _carries_marked_tool_call(message: Mapping[str, object]) -> bool:
|
||||
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
|
||||
tool_calls: Final = message.get("tool_calls")
|
||||
return isinstance(tool_calls, (list, tuple)) and any(
|
||||
_carries_cache_mark(_tool_call_fields(call)) for call in tool_calls
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -722,7 +722,7 @@ class TestPromptCacheBreakpointOnClaude:
|
|||
def test_a_30m_breakpoint_still_caches(
|
||||
self, client: PassthroughClient, resources: ResourceManager, scoped_key: str, route: ClaudeRoute
|
||||
) -> None:
|
||||
model: Final = _register_claude(client, resources, route)
|
||||
model: Final = _register_claude(client, resources, route, drop_params=True)
|
||||
body: Final = ChatBody(
|
||||
model=model,
|
||||
max_tokens=64,
|
||||
|
|
@ -760,7 +760,7 @@ class TestPromptCacheBreakpointOnClaude:
|
|||
def test_a_block_carrying_both_keys_still_caches(
|
||||
self, client: PassthroughClient, resources: ResourceManager, scoped_key: str, route: ClaudeRoute
|
||||
) -> None:
|
||||
model: Final = _register_claude(client, resources, route)
|
||||
model: Final = _register_claude(client, resources, route, drop_params=True)
|
||||
both_keys: Final = TextContentPart(
|
||||
text=_cacheable_prefix(),
|
||||
cache_control=CacheControl(),
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import Final, List, Optional, Tuple, cast
|
|||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
from openai.types.chat import ChatCompletionMessage
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
|
|
@ -1523,6 +1524,49 @@ def test_cache_control_hook_fills_the_mark_budget_every_claude_route_sends_after
|
|||
assert _marks_the_route_sends(model, custom_llm_provider, processed) == 4
|
||||
|
||||
|
||||
def test_a_breakpoint_on_a_copilot_tool_call_never_starts_the_github_device_login(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
|
||||
):
|
||||
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
|
||||
monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://github.test/login/device/code")
|
||||
messages: Final = cast(
|
||||
list[AllMessageValues],
|
||||
[
|
||||
{"role": "system", "content": "You are a weather assistant."},
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
"prompt_cache_breakpoint": {"mode": "explicit"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "Sunny"},
|
||||
],
|
||||
)
|
||||
|
||||
with respx.mock(assert_all_called=False) as github:
|
||||
device_login: Final = github.post("https://github.test/login/device/code").respond(status_code=503)
|
||||
AnthropicCacheControlHook().get_chat_completion_prompt(
|
||||
model="claude-sonnet-4.6",
|
||||
messages=messages,
|
||||
non_default_params={
|
||||
"custom_llm_provider": "github_copilot",
|
||||
"cache_control_injection_points": [{"location": "message", "role": "system"}],
|
||||
},
|
||||
prompt_id=None,
|
||||
prompt_variables=None,
|
||||
dynamic_callback_params={},
|
||||
)
|
||||
|
||||
assert not device_login.called
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool_call_type", ["function", "custom", None])
|
||||
@pytest.mark.parametrize(
|
||||
"mark",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
from functools import reduce
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -12,6 +13,7 @@ from litellm.router_utils.prompt_caching_cache import (
|
|||
PromptCachingCache,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message
|
||||
|
||||
|
||||
class _ManualClock:
|
||||
|
|
@ -321,6 +323,61 @@ def test_a_one_hour_breakpoint_on_a_tool_call_keeps_the_pin_after_the_mark_moves
|
|||
"model_id": "deployment-1h"
|
||||
}
|
||||
|
||||
|
||||
def _marked_lookup_call() -> ChatCompletionMessageToolCall:
|
||||
return ChatCompletionMessageToolCall(
|
||||
id="call_1",
|
||||
type="function",
|
||||
function=Function(name="lookup_clause", arguments="{}"),
|
||||
prompt_cache_breakpoint={"mode": "explicit", "ttl": "1h"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"marked_turn",
|
||||
[
|
||||
Message(content=None, tool_calls=[_marked_lookup_call()]),
|
||||
{"role": "assistant", "content": None, "tool_calls": [_marked_lookup_call()]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": ({**_tool_call("call_1"), "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}},),
|
||||
},
|
||||
],
|
||||
ids=["response_message", "response_tool_calls", "tuple_of_dicts"],
|
||||
)
|
||||
def test_a_one_hour_breakpoint_keeps_the_pin_whatever_shape_the_marked_tool_calls_come_in(marked_turn: object) -> None:
|
||||
clock: Final = _ManualClock()
|
||||
cache: Final = _pin_cache(clock)
|
||||
messages: Final = cast(list[AllMessageValues], [*_tool_call_turn(marked_call_id="call_1")[:-1], marked_turn])
|
||||
|
||||
cache.add_model_id(model_id="deployment-1h", messages=messages, tools=None, honors_one_hour_ttl=True)
|
||||
clock.advance(PROMPT_CACHE_PIN_TTL_SECONDS + 1)
|
||||
|
||||
assert cache.get_model_id(messages=messages, tools=None) == {"model_id": "deployment-1h"}
|
||||
|
||||
|
||||
def test_a_deeply_nested_value_on_a_tool_call_after_the_mark_keeps_the_pin() -> None:
|
||||
cache: Final = _pin_cache(_ManualClock())
|
||||
deep_call: Final = {**_tool_call("call_1"), "x": reduce(lambda inner, _: {"x": inner}, range(300), "leaf")}
|
||||
messages: Final = cast(
|
||||
list[AllMessageValues],
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "Here is the full text of a complex legal agreement" * 200,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{"role": "user", "content": "Find the termination clause."},
|
||||
{"role": "assistant", "content": None, "tool_calls": [deep_call]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "Section 9."},
|
||||
],
|
||||
)
|
||||
|
||||
cache.add_model_id(model_id="deployment-1", messages=messages, tools=None)
|
||||
|
||||
assert cache.get_model_id(messages=messages, tools=None) == {"model_id": "deployment-1"}
|
||||
|
||||
async def _first_non_none(read):
|
||||
while True:
|
||||
value = await read()
|
||||
|
|
|
|||
|
|
@ -1204,16 +1204,21 @@ def test_a_five_minute_mark_on_a_tool_the_route_sends_unmarked_leaves_the_system
|
|||
assert _sent_cache_ttls(json.loads(route.calls.last.request.content)) == ("1h", "1h")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"member",
|
||||
[
|
||||
{"type": "function", "name": "spawn", "parameters": {"type": "object"}, "cache_control": _FIVE_MINUTES},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": "spawn", "parameters": {"type": "object"}, "cache_control": _FIVE_MINUTES},
|
||||
},
|
||||
],
|
||||
ids=["flat", "function_form"],
|
||||
)
|
||||
def test_a_five_minute_mark_on_a_namespace_member_leaves_no_hour_after_it_under_drop_params(
|
||||
respx_mock: respx.MockRouter,
|
||||
respx_mock: respx.MockRouter, member: Mapping[str, object]
|
||||
):
|
||||
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").respond(json=_CLAUDE_REPLY)
|
||||
member: Final = {
|
||||
"type": "function",
|
||||
"name": "spawn",
|
||||
"parameters": {"type": "object"},
|
||||
"cache_control": _FIVE_MINUTES,
|
||||
}
|
||||
|
||||
litellm.completion(
|
||||
model="anthropic/claude-opus-5-5",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue