From 4735bae4b4311e3599fbbaad3bf0b31b16eda7fd Mon Sep 17 00:00:00 2001 From: sailikhithk Date: Fri, 14 Aug 2026 18:01:34 -0500 Subject: [PATCH] fix(vertex_ai): convert messages to contents in gemini count_tokens acount_tokens passed contents=None to the Vertex Gemini countTokens endpoint when called with messages=, causing a silent zero token count. The Gemini branch of VertexAITokenCounter.count_tokens never read the messages parameter, so the request body was {"contents": null}, which Vertex accepts with HTTP 200 and no totalTokens field. Convert messages to Gemini contents format using the existing _gemini_convert_messages_with_history helper when contents is None. Treat a response without totalTokens as a failure so the caller falls back to local token counting instead of returning a silent zero. Fixes #36921 --- litellm/llms/vertex_ai/common_utils.py | 16 +++- .../vertex_ai/test_vertex_ai_common_utils.py | 79 +++++++++++++++++++ 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 26f797cf5b2..e68470c6b49 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1164,21 +1164,29 @@ class VertexAITokenCounter(BaseTokenCounter): original_response=result, ) else: - # Use standard Vertex AI (Gemini) token counter from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter + from litellm.llms.vertex_ai.gemini.transformation import ( + _gemini_convert_messages_with_history, + ) + + resolved_contents: Final = ( + contents + if contents is not None + else _gemini_convert_messages_with_history(messages=messages or []) + ) count_tokens_params: Final = { "model": model_to_use, - "contents": contents, + "contents": resolved_contents, } count_tokens_params_request.update(count_tokens_params) result = await VertexAITokenCounter().acount_tokens( **count_tokens_params_request, ) - if result is not None: + if result is not None and "totalTokens" in result: return TokenCountResponse( - total_tokens=result.get("totalTokens", 0), + total_tokens=result["totalTokens"], request_model=request_model, model_used=model_to_use, tokenizer_type=result.get("tokenizer_used", ""), diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index b83d4742b64..7ebdf2d02c8 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1275,6 +1275,85 @@ async def test_vertex_ai_token_counter_routes_gemini_models(): assert result.total_tokens == 50 +@pytest.mark.asyncio +async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini(): + """ + Regression test for #36921: acount_tokens passed contents=None to the + Gemini countTokens endpoint when called with messages=, causing a + silent zero token count. Verify messages are converted to Gemini + contents format when contents is None. + """ + from unittest.mock import patch + + from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter + + token_counter = VertexAITokenCounter() + + with patch( + "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" + ) as mock_acount_tokens: + mock_acount_tokens.return_value = { + "totalTokens": 42, + "tokenizer_used": "gemini", + } + + await token_counter.count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello, how are you?"}], + contents=None, + deployment={ + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-central1", + } + }, + request_model="vertex_ai/gemini-2.5-flash", + ) + + mock_acount_tokens.assert_called_once() + call_kwargs = mock_acount_tokens.call_args.kwargs + passed_contents = call_kwargs["contents"] + assert passed_contents is not None + assert isinstance(passed_contents, list) + assert len(passed_contents) >= 1 + assert "parts" in passed_contents[0] + + +@pytest.mark.asyncio +async def test_vertex_ai_token_counter_returns_none_when_api_omits_total_tokens(): + """ + Regression test for #36921: Vertex returns HTTP 200 with no totalTokens + when contents is null. The old code read totalTokens with a default of 0 + and returned a silent zero. Verify we now return None so the caller falls + back to local token counting. + """ + from unittest.mock import patch + + from litellm.llms.vertex_ai.common_utils import VertexAITokenCounter + + token_counter = VertexAITokenCounter() + + with patch( + "litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens" + ) as mock_acount_tokens: + mock_acount_tokens.return_value = {"tokenizer_used": "gemini"} + + result = await token_counter.count_tokens( + model_to_use="gemini-2.5-flash", + messages=[{"role": "user", "content": "Hello"}], + contents=None, + deployment={ + "litellm_params": { + "vertex_project": "test-project", + "vertex_location": "us-central1", + } + }, + request_model="vertex_ai/gemini-2.5-flash", + ) + + assert result is None + + @pytest.mark.asyncio async def test_vertex_ai_partner_model_detection(): """