From 236701944793d8c6f291b9eeadcfbf894ee17d15 Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 28 Oct 2026 01:23:03 +0000 Subject: [PATCH 1/2] fix(proxy): restore legacy token count input precedence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 7 --- .../test_vertex_gemini_count_tokens_wire.py | 30 +++++-------- tests/unit/proxy/test_proxy_token_counter.py | 43 ++++--------------- 3 files changed, 19 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3b36ea5afce..76ef3e3b541 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15328,13 +15328,6 @@ async def count_request_tokens( sent_params: Final[Mapping[str, object]] = {**_router_default_litellm_params(llm_router), **request_litellm_params} count_deployment: Final = _deployment_with_request_drop_params(deployment, sent_params.get("drop_params")) conversation_empty: Final = not messages and not contents - if prompt and not conversation_empty and not deployment_drops_params(count_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 request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment) contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents)) 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 1632e424c55..5c08fddd4b8 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -1046,33 +1046,23 @@ def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_for (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)), ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), ) -def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider( - gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool +@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() with wire_server(peer) as wire, gateway.scenario() as scenario: - model: Final = _ai_studio_deployment(scenario, wire.url) - response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint) - requests: Final = _count_requests(wire) - assert response.status_code == 400, response.text - assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt" - assert requests == () - - -@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count")) -def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt( - gateway: Gateway, call_endpoint: bool -) -> None: - peer: Final = _GeminiPeer() - with wire_server(peer) as wire, gateway.scenario() as scenario: - model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True) - mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint) + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params) + mixed: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint) mixed_requests: Final = _count_requests(wire) - alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint) + alone: Final = _count(gateway, {"model": model, **conversation}, call_endpoint) alone_requests: Final = _count_requests(wire) assert mixed.status_code == 200, mixed.text assert alone.status_code == 200, alone.text - assert _total_tokens(mixed) == _total_tokens(alone) + mixed_tokens: Final = TypeAdapter(int).validate_python(_total_tokens(mixed)) + alone_tokens: Final = TypeAdapter(int).validate_python(_total_tokens(alone)) + 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 fa0c9eda02d..436694c3849 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3419,43 +3419,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, From 8d1489085fffe0b64fb9970356924f33a46b22ad Mon Sep 17 00:00:00 2001 From: joshua Date: Wed, 28 Oct 2026 01:53:06 +0000 Subject: [PATCH 2/2] fix(proxy): count a contents conversation locally instead of the prompt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 33 ++++++++++++++++++- .../test_vertex_gemini_count_tokens_wire.py | 9 +++-- tests/unit/proxy/test_proxy_token_counter.py | 9 +++-- 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 76ef3e3b541..450b66745d4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -350,6 +350,7 @@ from litellm.constants import ( WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError +from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting @@ -14762,6 +14763,24 @@ def _counted_request_messages( return list(drop_non_string_tool_result_texts(messages, deployment_drops_params(deployment))) +_LOCALLY_COUNTED_MESSAGES: Final = TypeAdapter(list[dict[str, object]]) + + +def _locally_counted_messages( + request_messages: list[dict[str, object]] | None, # mutable-ok: _counted_request_messages returns list messages + contents: Sequence[Mapping[str, object]] | None, + model: str, +) -> list[dict[str, object]] | None: # mutable-ok: token_counter takes list messages + if request_messages or not contents: + return request_messages + translate: Final = cast( # cast-ok: the adapter's untyped **kwargs leave its signature partially unknown + Callable[[str, list[dict[str, object]]], Mapping[str, object]], + GoogleGenAIAdapter().translate_generate_content_to_completion, + ) + translated: Final = translate(model, [dict(content) for content in contents]) + return _LOCALLY_COUNTED_MESSAGES.validate_python(translated.get("messages", [])) + + async def _try_provider_token_count( provider_counter: "BaseTokenCounter", custom_llm_provider: str | None, @@ -15413,7 +15432,19 @@ async def count_request_tokens( prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes Sequence[AllMessageValues] | None, - [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else request_messages), + [{"role": "user", "content": prompt}] + if prompt_turn + else ( + None + if prompt_only + else _locally_counted_messages( + request_messages, + cast( # cast-ok: TokenCountRequest.contents is typed list[dict], the local counter needs str-keyed parts + Sequence[Mapping[str, object]] | None, contents + ), + model_to_use, + ) + ), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) 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 5c08fddd4b8..214bee77895 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -1043,8 +1043,13 @@ def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_for @pytest.mark.parametrize( ("conversation", "call_endpoint"), - (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)), - ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), + ( + ({"messages": _MESSAGES}, True), + ({"contents": _CONTENTS}, True), + ({"messages": _MESSAGES}, False), + ({"contents": _CONTENTS}, False), + ), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages", "local-count-contents"), ) @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( diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 436694c3849..be7e50adf62 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -3422,8 +3422,13 @@ _CONVERSATIONS: Final = ( @pytest.mark.asyncio @pytest.mark.parametrize( ("conversation", "call_endpoint"), - ((_CONVERSATIONS[0], True), (_CONVERSATIONS[1], True), (_CONVERSATIONS[0], False)), - ids=("provider_messages", "provider_contents", "local_messages"), + ( + (_CONVERSATIONS[0], True), + (_CONVERSATIONS[1], True), + (_CONVERSATIONS[0], False), + (_CONVERSATIONS[1], False), + ), + ids=("provider_messages", "provider_contents", "local_messages", "local_contents"), ) @pytest.mark.parametrize( ("global_drop_params", "deployment_drop_params"),