Merge pull request #14133 from retanoj/fix/gemini-token-count

Fix token count error for gemini cli

Author: krrishdholakia

litellm/google_genai/adapters/handler.py

@@ -37,6 +37,10 @@ def _prepare_completion_kwargs(
 
         completion_kwargs: Dict[str, Any] = dict(completion_request)
 
+        # feed metadata for custom callback
+        if extra_kwargs is not None and "metadata" in extra_kwargs:
+            completion_kwargs["metadata"] = extra_kwargs["metadata"]
+
         if stream:
             completion_kwargs["stream"] = stream
 

litellm/proxy/google_endpoints/endpoints.py

@@ -173,15 +173,24 @@ 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,11 +201,17 @@ 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
     #########################################################

tests/test_litellm/proxy/google_endpoints/__init__.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