mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
Fix tests
This commit is contained in:
parent
1284e4ebe5
commit
6146196c6a
1 changed files with 15 additions and 5 deletions
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue