diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py b/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py deleted file mode 100644 index c8894ca7d5..0000000000 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_gemini_token_usage.py +++ /dev/null @@ -1,70 +0,0 @@ - -import sys, os -import pytest -sys.path.insert(0, os.path.abspath('../../../../../')) - -from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig -from litellm.types.llms.vertex_ai import UsageMetadata - -def test_gemini_3_flash_preview_token_usage_fallback(): - """Test fallback logic when candidatesTokensDetails is missing (e.g. Gemini 3 Flash Preview).""" - v = VertexGeminiConfig() - - usage_metadata_dict = { - "promptTokenCount": 2145, - "candidatesTokenCount": 509, - "totalTokenCount": 2654, - # candidatesTokensDetails intentionally omitted - } - - completion_response = {"usageMetadata": usage_metadata_dict} - result = v._calculate_usage(completion_response=completion_response) - - assert result.completion_tokens == 509 - assert result.prompt_tokens == 2145 - assert result.total_tokens == 2654 - - # Text tokens should be derived from candidatesTokenCount - assert result.completion_tokens_details is not None - assert result.completion_tokens_details.text_tokens == 509 - assert result.completion_tokens_details.image_tokens is None - assert result.completion_tokens_details.audio_tokens is None - -def test_gemini_no_reasoning_fallback(): - """Test fallback when reasoning_effort is absent and details are missing.""" - v = VertexGeminiConfig() - - usage_metadata_dict = { - "promptTokenCount": 100, - "candidatesTokenCount": 264, - "totalTokenCount": 364, - } - - completion_response = {"usageMetadata": usage_metadata_dict} - result = v._calculate_usage(completion_response=completion_response) - - assert result.completion_tokens == 264 - assert result.completion_tokens_details is not None - assert result.completion_tokens_details.text_tokens == 264 - assert result.completion_tokens_details.reasoning_tokens is None or result.completion_tokens_details.reasoning_tokens == 0 - -def test_gemini_token_usage_standard_response(): - """Verify that standard responses with details are computed correctly and not overwritten.""" - v = VertexGeminiConfig() - - usage_metadata_dict = { - "promptTokenCount": 100, - "candidatesTokenCount": 50, - "totalTokenCount": 150, - "candidatesTokensDetails": [ - {"modality": "TEXT", "tokenCount": 40}, - {"modality": "IMAGE", "tokenCount": 10} - ] - } - - completion_response = {"usageMetadata": usage_metadata_dict} - result = v._calculate_usage(completion_response=completion_response) - - assert result.completion_tokens == 50 - assert result.completion_tokens_details.text_tokens == 40 - assert result.completion_tokens_details.image_tokens == 10 diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index d09de3a0f2..5d91275fa8 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2505,3 +2505,72 @@ def test_vertex_ai_multiple_function_declarations_grouped(): func_names = [f["name"] for f in tools[0]["function_declarations"]] assert "func1" in func_names assert "func2" in func_names + + +def test_gemini_3_flash_preview_token_usage_fallback(): + """Test fallback logic when candidatesTokensDetails is missing (e.g. Gemini 3 Flash Preview).""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 2145, + "candidatesTokenCount": 509, + "totalTokenCount": 2654, + # candidatesTokensDetails intentionally omitted + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 509 + assert result.prompt_tokens == 2145 + assert result.total_tokens == 2654 + + # Text tokens should be derived from candidatesTokenCount + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 509 + assert result.completion_tokens_details.image_tokens is None + assert result.completion_tokens_details.audio_tokens is None + + +def test_gemini_no_reasoning_fallback(): + """Test fallback when reasoning_effort is absent and details are missing.""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 100, + "candidatesTokenCount": 264, + "totalTokenCount": 364, + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 264 + assert result.completion_tokens_details is not None + assert result.completion_tokens_details.text_tokens == 264 + assert ( + result.completion_tokens_details.reasoning_tokens is None + or result.completion_tokens_details.reasoning_tokens == 0 + ) + + +def test_gemini_token_usage_standard_response(): + """Verify that standard responses with details are computed correctly and not overwritten.""" + v = VertexGeminiConfig() + + usage_metadata_dict = { + "promptTokenCount": 100, + "candidatesTokenCount": 50, + "totalTokenCount": 150, + "candidatesTokensDetails": [ + {"modality": "TEXT", "tokenCount": 40}, + {"modality": "IMAGE", "tokenCount": 10}, + ], + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + assert result.completion_tokens == 50 + assert result.completion_tokens_details.text_tokens == 40 + assert result.completion_tokens_details.image_tokens == 10