scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(router): count tools and Anthropic system prompt in context-window pre-call check (#39663)

Browse source 
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-03 16:38:48 -07:00 • committed by GitHub
parent ab0478f068
commit 0b7773dd44
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 198 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/anthropic/experimental_pass_through/messages/utils.py
View file

@ -6,6 +6,7 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.llms.openai import ChatCompletionSystemMessage
if TYPE_CHECKING:
from litellm.exceptions import ContentPolicyViolationError
@ -36,6 +37,16 @@ def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "
)
def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMessage | None:
"""
Return the Anthropic Messages top-level ``system`` (a string or a list of text
blocks) as an OpenAI-style system message, or None when the request has none.
"""
if not isinstance(system, (str, list)) or not system:
return None
return ChatCompletionSystemMessage(role="system", content=system)
@lru_cache(maxsize=1)
def _anthropic_messages_optional_param_keys() -> frozenset[str]:
"""

36
litellm/router.py
View file

@ -197,6 +197,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import (
from litellm.scheduler import FlowItem, Scheduler
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolParam,
FileTypes,
OpenAIFileObject,
OpenAIFilesPurpose,
@ -11762,7 +11763,7 @@ class Router:
self,
messages: list[dict[str, str]] | None,
input: str | list | None,
instructions: str | None = None,
request_kwargs: Mapping[str, object] | None = None,
) -> int:
"""
Count input tokens for context-window pre-call checks.
@ -11772,9 +11773,28 @@ class Router:
The Responses payload is normalized to chat messages via the shared
LiteLLMCompletionResponsesConfig transform so the same token_counter path covers
both API surfaces and `instructions` tokens are included in the count.
Prompt content the message list never carries is read from `request_kwargs`:
`tools` (Chat Completions, Responses and Anthropic Messages shapes) and the
Anthropic Messages top-level `system` block.
"""
from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
anthropic_system_to_openai_message,
)
extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({})
raw_instructions: Final = extras.get("instructions")
instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
raw_tools: Final = extras.get("tools")
tools: Final = (
cast(list[ChatCompletionToolParam], raw_tools) # cast-ok: token_counter formats any tool dict shape
if isinstance(raw_tools, list) and raw_tools
else None
)
system_message: Final = anthropic_system_to_openai_message(extras.get("system"))
if messages is not None:
return litellm.token_counter(messages=messages)
counted_messages: Final = (system_message, *messages) if system_message is not None else messages
return litellm.token_counter(messages=counted_messages, tools=tools)
if input is not None:
from openai.types.responses.response_create_params import ResponseInputParam
@ -11787,7 +11807,10 @@ class Router:
input=typed_input,
responses_api_request={"instructions": instructions} if instructions is not None else {},
)
return litellm.token_counter(messages=cast(list, input_messages)) # cast-ok: transformed chat messages
return litellm.token_counter(
messages=cast(list, input_messages), # cast-ok: transformed chat messages
tools=tools,
)
raise ValueError("Either messages or input must be provided to count tokens")
def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None:
@ -11833,14 +11856,13 @@ class Router:
"""
if messages is None and input is None:
return None
raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
try:
if not self._pre_call_checks_need_token_count(model, healthy_deployments):
return None
return await asyncify(self._count_pre_call_check_tokens)(
messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter
input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter
instructions=raw_instructions if isinstance(raw_instructions, str) else None,
request_kwargs=request_kwargs,
)
except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request
verbose_router_logger.error(
@ -11887,8 +11909,6 @@ class Router:
_rate_limit_error = False
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs)
raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
has_countable_input: Final = messages is not None or input is not None
## get model group RPM ##
@ -11919,7 +11939,7 @@ class Router:
return _returned_deployments
try:
input_tokens = self._count_pre_call_check_tokens(
messages=messages, input=input, instructions=instructions
messages=messages, input=input, request_kwargs=request_kwargs
)
except Exception as e:
verbose_router_logger.error(

160
tests/test_litellm/test_router.py
View file

@ -3855,7 +3855,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch):
input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input)
with_instructions_tokens = router._count_pre_call_check_tokens(
messages=None, input=short_input, instructions=long_instructions
messages=None, input=short_input, request_kwargs={"instructions": long_instructions}
)
assert with_instructions_tokens > input_only_tokens
@ -3871,6 +3871,164 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch):
)
_OVERSIZED_TOOL_DESCRIPTION = "look up the answer in the knowledge base. " * 40
@pytest.mark.parametrize(
"prompt_kwargs, tool",
[
pytest.param(
{"messages": [{"role": "user", "content": "hi"}]},
{
"type": "function",
"function": {
"name": "lookup",
"description": _OVERSIZED_TOOL_DESCRIPTION,
"parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
},
},
id="chat_completions_tool",
),
pytest.param(
{"input": "hi"},
{
"type": "function",
"name": "lookup",
"description": _OVERSIZED_TOOL_DESCRIPTION,
"parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
},
id="responses_tool",
),
pytest.param(
{"messages": [{"role": "user", "content": "hi"}]},
{
"name": "lookup",
"description": _OVERSIZED_TOOL_DESCRIPTION,
"input_schema": {"type": "object", "properties": {"q": {"type": "string"}}},
},
id="anthropic_messages_tool",
),
],
)
def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwargs, tool):
"""
Tool definitions are sent to the model as prompt tokens but never appear in
`messages` or `input`. A request whose prompt alone fits the context window but
whose prompt plus `tools` exceeds it must be rejected before dispatch, for the
Chat Completions, Responses and Anthropic Messages tool shapes alike.
"""
router = litellm.Router(
model_list=[
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
],
enable_pre_call_checks=True,
)
deployments = [
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
]
prompt_only_tokens = router._count_pre_call_check_tokens(
messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input")
)
monkeypatch.setattr(
router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}
)
assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1
with pytest.raises(litellm.ContextWindowExceededError):
router._pre_call_checks(
model="m",
healthy_deployments=deployments,
request_kwargs={"tools": [tool]},
**prompt_kwargs,
)
@pytest.mark.parametrize(
"system",
[
pytest.param("You are a meticulous assistant. " * 40, id="system_string"),
pytest.param(
[{"type": "text", "text": "You are a meticulous assistant. " * 40}],
id="system_blocks",
),
],
)
def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system):
"""
The Anthropic Messages API carries the system prompt as a top-level `system` field,
not as a message. Its tokens reach the model, so a request whose `messages` fit but
whose `messages` plus `system` exceed the context window must be rejected.
"""
router = litellm.Router(
model_list=[
{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
],
enable_pre_call_checks=True,
)
deployments = [
{"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
]
messages = [{"role": "user", "content": "hi"}]
messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None)
monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens})
assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1
with pytest.raises(litellm.ContextWindowExceededError):
router._pre_call_checks(
model="m",
healthy_deployments=deployments,
messages=messages,
request_kwargs={"system": system},
)
@pytest.mark.asyncio
async def test_aanthropic_messages_enforces_context_window_with_system_and_tools():
"""
End-to-end router regression for /v1/messages: a request whose only oversized
content lives in the top-level `system` field or in `tools` must trip the pre-call
context-window check instead of being dispatched (the deployment uses mock_response,
so reaching the provider handler would return a response rather than raise).
"""
router = litellm.Router(
model_list=[
{
"model_name": "small-ctx",
"litellm_params": {"model": "anthropic/claude-3-5-haiku-20241022", "mock_response": "hi"},
"model_info": {"max_input_tokens": 20},
}
],
enable_pre_call_checks=True,
)
messages = [{"role": "user", "content": "hi"}]
response = await router.aanthropic_messages(model="small-ctx", messages=messages, max_tokens=5)
assert response is not None
with pytest.raises(litellm.ContextWindowExceededError):
await router.aanthropic_messages(
model="small-ctx",
messages=messages,
max_tokens=5,
system="You are a meticulous assistant. " * 40,
)
with pytest.raises(litellm.ContextWindowExceededError):
await router.aanthropic_messages(
model="small-ctx",
messages=messages,
max_tokens=5,
tools=[
{
"name": "lookup",
"description": _OVERSIZED_TOOL_DESCRIPTION,
"input_schema": {"type": "object", "properties": {"q": {"type": "string"}}},
}
],
)
def test_count_pre_call_check_tokens_across_api_surfaces():
"""
_count_pre_call_check_tokens must count tokens from chat `messages`, a Responses