diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index da76b12c37..ae1a758bf2 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -65,6 +65,8 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): # Use custom api_base if provided, otherwise construct default if api_base: base_url = api_base + elif vertex_location == "global": + base_url = "https://aiplatform.googleapis.com" else: base_url = f"https://{vertex_location}-aiplatform.googleapis.com" diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 39606ef7ce..eb364ebf5b 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -819,3 +819,20 @@ async def test_vertex_ai_anthropic_token_counting(): assert response.original_response is not None assert "input_tokens" in response.original_response assert response.original_response["input_tokens"] == 15 + +@pytest.mark.parametrize("vertex_location", ["global", "us-central1"]) +def test_vertex_ai_gemini_token_counting_endpoint(vertex_location): + from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import ( + VertexAIPartnerModelsTokenCounter, + ) + + endpoint = VertexAIPartnerModelsTokenCounter()._build_count_tokens_endpoint( + model="gemini-2.5-pro", + project_id="test-project", + vertex_location=vertex_location, + api_base=None, + ) + if vertex_location == "global": + assert endpoint == "https://aiplatform.googleapis.com" + else: + assert endpoint == f"https://{vertex_location}-aiplatform.googleapis.com" \ No newline at end of file