From 29b219ae2140185d313e3bb9b779b43d753d8fbd Mon Sep 17 00:00:00 2001
From: Atharva Jaiswal <92455570+AtharvaJaiswal005@users.noreply.github.com>
Date: Thu, 19 Mar 2026 12:33:32 +0530
Subject: [PATCH] fix(vertex): respect vertex_count_tokens_location for Claude
 models

handler.py unconditionally overrides vertex_location to us-central1
when the model name contains "claude", ignoring any explicitly configured
vertex_count_tokens_location. Check for vertex_count_tokens_location
first and only fall back to us-central1 when it is not set.

Fixes #23872
---
 .../count_tokens/handler.py                   | 30 ++++++++-----------
 1 file changed, 12 insertions(+), 18 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 5d94cd42129..93797e0a349 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
@@ -105,27 +105,21 @@ class VertexAIPartnerModelsTokenCounter(VertexBase):
         # Extract Vertex AI credentials and settings
         vertex_credentials = self.get_vertex_ai_credentials(litellm_params)
         vertex_project = self.get_vertex_ai_project(litellm_params)
+        vertex_location = self.get_vertex_ai_location(litellm_params)
 
-        # Check for count_tokens specific location override
-        vertex_count_tokens_location = litellm_params.get(
+        # 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.
+        # 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"
         )
-        vertex_location_raw = self.get_vertex_ai_location(litellm_params)
-
-        # Determine final location with precedence:
-        # 1. vertex_count_tokens_location (if provided)
-        # 2. vertex_location (if provided)
-        # 3. Default to us-east5 for Claude models when no location is set
-        # Supported regions: us-east5, europe-west1, asia-southeast1
-        # https://docs.cloud.google.com/vertex-ai/generative-ai/docs/partner-models/claude/count-tokens
-        if vertex_count_tokens_location:
-            vertex_location: str = vertex_count_tokens_location
-        elif vertex_location_raw:
-            vertex_location = vertex_location_raw
-        elif "claude" in model.lower():
-            vertex_location = "us-east5"
-        else:
-            vertex_location = "us-east5"
+        if explicit_count_tokens_location:
+            vertex_location = explicit_count_tokens_location
+        elif not vertex_location or "claude" in model.lower():
+            vertex_location = "us-central1"
 
         # Get access token and resolved project ID
         access_token, project_id = await self._ensure_access_token_async(