fix(proxy): restore legacy token count input precedence

This commit is contained in:
Joshua Valluru 2026-10-19 17:28:56 -07:00
parent a63df167da
commit e109f8977d
3 changed files with 45 additions and 48 deletions

View file

@ -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 = (

View file

@ -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
]

View file

@ -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,