From e894e7dab36bbfce38dc3bc165d22d3f8e75ced0 Mon Sep 17 00:00:00 2001 From: Atharva Jaiswal <92455570+AtharvaJaiswal005@users.noreply.github.com> Date: Wed, 8 Apr 2026 16:47:57 +0530 Subject: [PATCH] fix(vertex): respect existing vertex_location, default Claude to us-east5 Address upstream test expectations in test_count_tokens_location.py: - vertex_count_tokens_location override takes precedence (unchanged) - explicit vertex_location is now respected even for Claude models - Claude with no location set defaults to us-east5 (was us-central1) The previous fallback was too aggressive: it overrode any vertex_location to us-central1 for Claude requests, breaking users who explicitly set a supported region like asia-southeast1. --- .../count_tokens/handler.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) 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 de08f1fc587..6188c7c8c6e 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 @@ -108,19 +108,20 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): vertex_project = self.get_vertex_ai_project(litellm_params) vertex_location = self.get_vertex_ai_location(litellm_params) - # common_utils.py already resolves vertex_count_tokens_location and - # passes it as vertex_location through the call chain. - # Only override to us-central1 when the user did NOT explicitly set - # vertex_count_tokens_location -- this preserves backwards compat for - # users whose vertex_location is "global" or another unsupported region. + # Resolve the location used for the count-tokens endpoint with the + # following precedence: + # 1. vertex_count_tokens_location (explicit override) wins + # 2. vertex_location (if explicitly set) is respected + # 3. Claude models with no location set default to us-east5 + # Supported Claude regions: us-east5, europe-west1, asia-southeast1 # https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens explicit_count_tokens_location = litellm_params.get( "vertex_count_tokens_location" ) if explicit_count_tokens_location: vertex_location = explicit_count_tokens_location - elif not vertex_location or "claude" in model.lower(): - vertex_location = "us-central1" + elif not vertex_location and "claude" in model.lower(): + vertex_location = "us-east5" # Get access token and resolved project ID access_token, project_id = await self._ensure_access_token_async(