diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 1f575f2759..dcf707ebd5 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 extra_kwargs is not None and "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..eb481b0a4f 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/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 ######################################################### 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