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 079a691395a..ceb924b9b0f 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,16 +105,15 @@ 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 = ( - litellm_params.get("vertex_count_tokens_location") - or self.get_vertex_ai_location(litellm_params) - ) + 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 # 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 and "claude" in model.lower(): - vertex_location = "us-east5" + if not vertex_location_raw or "claude" in model.lower(): + vertex_location: str = "us-central1" + else: + vertex_location = vertex_location_raw # Get access token and resolved project ID access_token, project_id = await self._ensure_access_token_async( diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 2b71b88356b..29265bb4b42 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -122,21 +122,21 @@ def test_add_transform_inline_image_block_skips_data_urls(): # str branch str_content = {"type": "image_url", "image_url": data_url} result = config._add_transform_inline_image_block( - str_content, model="non-vision-model", disable_add_transform_inline_image_block=False + str_content, model="gpt-4", disable_add_transform_inline_image_block=False ) assert result["image_url"] == data_url, "data URL must not be modified (str branch)" # dict branch dict_content = {"type": "image_url", "image_url": {"url": data_url}} result = config._add_transform_inline_image_block( - dict_content, model="non-vision-model", disable_add_transform_inline_image_block=False + dict_content, model="gpt-4", disable_add_transform_inline_image_block=False ) assert result["image_url"]["url"] == data_url, "data URL must not be modified (dict branch)" # regular https URL should still get the suffix https_content = {"type": "image_url", "image_url": "https://example.com/image.jpg"} result = config._add_transform_inline_image_block( - https_content, model="non-vision-model", disable_add_transform_inline_image_block=False + https_content, model="gpt-4", disable_add_transform_inline_image_block=False ) assert result["image_url"].endswith("#transform=inline"), "https URL should get #transform=inline" @pytest.mark.parametrize(