From 002c2f16dc679c34944dfb9f68dc7595dc1ab9a9 Mon Sep 17 00:00:00 2001 From: retanoj Date: Mon, 1 Sep 2025 20:02:53 +0800 Subject: [PATCH 1/4] fix token count error when proxy gemini cli to openai like model --- litellm/google_genai/adapters/handler.py | 4 ++ litellm/proxy/google_endpoints/endpoints.py | 23 +++++++-- .../proxy/google_endpoints/__init__.py | 0 .../proxy/google_endpoints/test_endpoints.py | 49 +++++++++++++++++++ 4 files changed, 71 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/google_endpoints/__init__.py create mode 100644 tests/test_litellm/proxy/google_endpoints/test_endpoints.py diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 1f575f2759..c5f378554b 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -37,6 +37,10 @@ class GenerateContentToCompletionHandler: completion_kwargs: Dict[str, Any] = dict(completion_request) + # feed metadata for custom callback + # if 'metadata' in extra_kwargs: + # completion_kwargs['metadata'] = extra_kwargs['metadata'] + if stream: completion_kwargs["stream"] = stream diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 373232e22d..4f57e1e7ce 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -173,15 +173,22 @@ async def google_count_tokens(request: Request, model_name: str): """ from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.proxy_server import token_counter as internal_token_counter + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter data = await _read_request_body(request=request) contents = data.get("contents", []) #Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest + # Translate contents to openai format messages using the adapter + messages = (GoogleGenAIAdapter() + .translate_generate_content_to_completion(model_name, contents) + .get("messages", [])) + token_request = TokenCountRequest( model=model_name, - contents=contents + contents=contents, + messages=messages, # compatibility when use openai-like endpoint ) # Call the internal token counter function with direct request flag set to False @@ -192,10 +199,16 @@ async def google_count_tokens(request: Request, model_name: str): if token_response is not None: # cast the response to the well known format original_response: dict = token_response.original_response or {} - return TokenCountDetailsResponse( - totalTokens=original_response.get("totalTokens", 0), - promptTokensDetails=original_response.get("promptTokensDetails", []), - ) + if original_response: + return TokenCountDetailsResponse( + totalTokens=original_response.get("totalTokens", 0), + promptTokensDetails=original_response.get("promptTokensDetails", []), + ) + else: + return TokenCountDetailsResponse( + totalTokens=token_response.total_tokens or 0, + promptTokensDetails=[], + ) ######################################################### # Return the response in the well known format diff --git a/tests/test_litellm/proxy/google_endpoints/__init__.py b/tests/test_litellm/proxy/google_endpoints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/proxy/google_endpoints/test_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py new file mode 100644 index 0000000000..2f2538bf9a --- /dev/null +++ b/tests/test_litellm/proxy/google_endpoints/test_endpoints.py @@ -0,0 +1,49 @@ +""" +Test for google_endpoints/endpoints.py +""" +import pytest +import sys, os +from dotenv import load_dotenv + + +from litellm.proxy.google_endpoints.endpoints import google_count_tokens +from litellm.types.llms.vertex_ai import TokenCountDetailsResponse +from starlette.requests import Request + +load_dotenv() + +sys.path.insert( + 0, os.path.abspath("../../../..") +) + +@pytest.mark.asyncio +async def test_proxy_gemini_to_openai_like_model_token_counting(): + """ + Test the token counting endpoint for proxing gemini to openai-like models. + """ + response: TokenCountDetailsResponse = await google_count_tokens( + request=Request( + scope={ + "type": "http", + "parsed_body": ( + [ + "contents" + ], + { + "contents": [ + { + "parts": [ + { + "text": "Hello, how are you?" + } + ] + } + ] + } + ) + } + ), + model_name="volcengine/foo", + ) + + assert response.get("totalTokens") > 0 \ No newline at end of file From e05ffcfb4c2b3bbcb9de75a8f1609ffa141f6326 Mon Sep 17 00:00:00 2001 From: retanoj Date: Mon, 1 Sep 2025 20:09:19 +0800 Subject: [PATCH 2/4] fix back --- litellm/google_genai/adapters/handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index c5f378554b..c15d0cb9de 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -38,8 +38,8 @@ class GenerateContentToCompletionHandler: completion_kwargs: Dict[str, Any] = dict(completion_request) # feed metadata for custom callback - # if 'metadata' in extra_kwargs: - # completion_kwargs['metadata'] = extra_kwargs['metadata'] + if 'metadata' in extra_kwargs: + completion_kwargs['metadata'] = extra_kwargs['metadata'] if stream: completion_kwargs["stream"] = stream From 3b524ba5f1084d5efc68dcc2b73f37160d37eace Mon Sep 17 00:00:00 2001 From: retanoj Date: Tue, 2 Sep 2025 09:52:50 +0800 Subject: [PATCH 3/4] format --- litellm/google_genai/adapters/handler.py | 4 ++-- litellm/proxy/google_endpoints/endpoints.py | 10 ++++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index c15d0cb9de..ee9083cbbf 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -38,8 +38,8 @@ class GenerateContentToCompletionHandler: completion_kwargs: Dict[str, Any] = dict(completion_request) # feed metadata for custom callback - if 'metadata' in extra_kwargs: - completion_kwargs['metadata'] = extra_kwargs['metadata'] + if "metadata" in extra_kwargs: + completion_kwargs["metadata"] = extra_kwargs["metadata"] if stream: completion_kwargs["stream"] = stream diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 4f57e1e7ce..eb481b0a4f 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -181,9 +181,11 @@ async def google_count_tokens(request: Request, model_name: str): from litellm.proxy._types import TokenCountRequest # Translate contents to openai format messages using the adapter - messages = (GoogleGenAIAdapter() - .translate_generate_content_to_completion(model_name, contents) - .get("messages", [])) + messages = ( + GoogleGenAIAdapter() + .translate_generate_content_to_completion(model_name, contents) + .get("messages", []) + ) token_request = TokenCountRequest( model=model_name, @@ -209,7 +211,7 @@ async def google_count_tokens(request: Request, model_name: str): totalTokens=token_response.total_tokens or 0, promptTokensDetails=[], ) - + ######################################################### # Return the response in the well known format ######################################################### From b1686ecbe5fd6d6bc8590cbe1294389548ff572b Mon Sep 17 00:00:00 2001 From: retanoj Date: Tue, 2 Sep 2025 09:59:31 +0800 Subject: [PATCH 4/4] format --- litellm/google_genai/adapters/handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index ee9083cbbf..dcf707ebd5 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -38,7 +38,7 @@ class GenerateContentToCompletionHandler: completion_kwargs: Dict[str, Any] = dict(completion_request) # feed metadata for custom callback - if "metadata" in extra_kwargs: + if extra_kwargs is not None and "metadata" in extra_kwargs: completion_kwargs["metadata"] = extra_kwargs["metadata"] if stream: