From 6146196c6ad938bdfaab41cb20b9c6a90f36b5a4 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 19 Mar 2026 16:43:39 +0530 Subject: [PATCH] Fix tests --- .../count_tokens/handler.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 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 ceb924b9b0f..82076ff3600 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,15 +105,25 @@ 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) + + # Check for count_tokens specific location override + vertex_count_tokens_location = litellm_params.get("vertex_count_tokens_location") vertex_location_raw = self.get_vertex_ai_location(litellm_params) - - # Default Claude models to us-east5 for count-tokens endpoint when no location is set + + # 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 not vertex_location_raw or "claude" in model.lower(): - vertex_location: str = "us-central1" - else: + 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" # Get access token and resolved project ID access_token, project_id = await self._ensure_access_token_async(