From e109f8977ddf3b266d2d6e9376be47374303f479 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:28:56 -0700 Subject: [PATCH] fix(proxy): restore legacy token count input precedence --- litellm/proxy/proxy_server.py | 12 ------ .../test_vertex_gemini_count_tokens_wire.py | 38 +++++++++++++++- tests/unit/proxy/test_proxy_token_counter.py | 43 ++++--------------- 3 files changed, 45 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 63a920628b6..79b4b8aaa29 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13996,11 +13996,6 @@ _PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( ) -def _drops_params(deployment: Mapping[str, object] | None) -> bool: - litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) - return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) - - def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14088,13 +14083,6 @@ async def count_request_tokens( model_to_use = _model conversation_empty: Final = not messages and not contents - if prompt and not conversation_empty and not _drops_params(deployment): - raise ProxyException( - message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", - type="invalid_request_error", - param="prompt", - code=400, - ) prompt_only: Final = bool(prompt) and conversation_empty prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES ladder: Final = ( diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 5fe82d37e9e..e54fb1a4a13 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -220,8 +220,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -594,3 +594,37 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + ( + ({"messages": _MESSAGES}, True), + ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]}, True), + ({"messages": _MESSAGES}, False), + ), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +@pytest.mark.parametrize("drop_params", (False, True), ids=("no-drop", "drop")) +def test_utils_token_counter_counts_the_conversation_regardless_of_prompt_and_drop_params( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool, drop_params: bool +) -> None: + peer: Final = _GeminiPeer() + params: Final = {"call_endpoint": "true" if call_endpoint else "false"} + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params) + mixed: Final = gateway.request( + "POST", "/utils/token_counter", {"model": model, "prompt": "hi", **conversation}, params=params + ) + mixed_requests: Final = _count_requests(wire) + alone: Final = gateway.request("POST", "/utils/token_counter", {"model": model, **conversation}, params=params) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + mixed_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(mixed.content)["total_tokens"]) + alone_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(alone.content)["total_tokens"]) + assert mixed_tokens == alone_tokens > 0 + assert len(mixed_requests) == len(alone_requests) == int(call_endpoint) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 995089105c4..9703e8ad1dc 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2515,43 +2515,18 @@ _CONVERSATIONS: Final = ( ) -@pytest.mark.asyncio -@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) -@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) -async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( - conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) - client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) - - with ( - patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), - pytest.raises(ProxyException) as rejected, - ): - await token_counter( - request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint - ) - - assert (rejected.value.code, rejected.value.param) == ("400", "prompt") - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), - [ - (_CONVERSATIONS[0], True, False, True), - (_CONVERSATIONS[1], True, False, True), - (_CONVERSATIONS[0], False, False, True), - (_CONVERSATIONS[0], False, True, False), - ], - ids=[ - "provider_messages", - "provider_contents", - "local_messages", - "local_messages_global_setting", - ], + ("conversation", "call_endpoint"), + ((_CONVERSATIONS[0], True), (_CONVERSATIONS[1], True), (_CONVERSATIONS[0], False)), + ids=("provider_messages", "provider_contents", "local_messages"), ) -async def test_drop_params_counts_a_conversation_without_the_prompt( +@pytest.mark.parametrize( + ("global_drop_params", "deployment_drop_params"), + ((False, False), (False, True), (True, False), (True, True)), + ids=("no_drop", "deployment_drop", "global_drop", "both_drop"), +) +async def test_conversation_takes_precedence_over_prompt_regardless_of_drop_params( conversation: dict[str, JsonValue], call_endpoint: bool, global_drop_params: bool,