diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index c995583e53..23facabbf8 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -496,3 +496,23 @@ def construct_target_url( updated_url = new_base_url.copy_with(path=updated_requested_route) return updated_url + + +def is_global_only_vertex_model(model: str) -> bool: + """ + Check if a model is only available in the global region. + + Args: + model: The model name to check + + Returns: + True if the model is only available in global region, False otherwise + """ + from litellm.utils import get_supported_regions + + supported_regions = get_supported_regions( + model=model, custom_llm_provider="vertex_ai" + ) + if supported_regions is None: + return False + return "global" in supported_regions diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 261b0876ff..5ea443513b 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -13,7 +13,12 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES, VertexPartnerProvider -from .common_utils import _get_gemini_url, _get_vertex_url, all_gemini_url_modes +from .common_utils import ( + _get_gemini_url, + _get_vertex_url, + all_gemini_url_modes, + is_global_only_vertex_model, +) if TYPE_CHECKING: from google.auth.credentials import Credentials as GoogleCredentialsObject @@ -34,7 +39,9 @@ class VertexBase: self.project_id: Optional[str] = None self.async_handler: Optional[AsyncHTTPHandler] = None - def get_vertex_region(self, vertex_region: Optional[str]) -> str: + def get_vertex_region(self, vertex_region: Optional[str], model: str) -> str: + if is_global_only_vertex_model(model): + return "global" return vertex_region or "us-central1" def load_auth( @@ -323,7 +330,10 @@ class VertexBase: ) auth_header = None # this field is not used for gemin else: - vertex_location = self.get_vertex_region(vertex_region=vertex_location) + vertex_location = self.get_vertex_region( + vertex_region=vertex_location, + model=model, + ) ### SET RUNTIME ENDPOINT ### version: Literal["v1beta1", "v1"] = ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 888097582d..3cda936572 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23,6 +23,13 @@ "search_context_size_medium": 0.0, "search_context_size_high": 0.0 }, + "supported_regions": [ + "global", + "us-west-2", + "eu-west-1", + "ap-southeast-1", + "ap-northeast-1" + ], "deprecation_date": "date when the model becomes deprecated in the format YYYY-MM-DD" }, "omni-moderation-latest": { @@ -6860,6 +6867,9 @@ "supported_output_modalities": [ "text" ], + "supported_regions": [ + "global" + ], "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", "supports_parallel_function_calling": true, "supports_web_search": true diff --git a/litellm/utils.py b/litellm/utils.py index 37fc3f2ff8..5d4effcf81 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -532,9 +532,9 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] - ] = kwargs.pop("callbacks", None) + dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + kwargs.pop("callbacks", None) + ) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) if len(all_callbacks) > 0: @@ -1240,9 +1240,9 @@ def client(original_function): # noqa: PLR0915 exception=e, retry_policy=kwargs.get("retry_policy"), ) - kwargs[ - "retry_policy" - ] = reset_retry_policy() # prevent infinite loops + kwargs["retry_policy"] = ( + reset_retry_policy() + ) # prevent infinite loops litellm.num_retries = ( None # set retries to None to prevent infinite loops ) @@ -2077,6 +2077,43 @@ def supports_reasoning(model: str, custom_llm_provider: Optional[str] = None) -> ) +def get_supported_regions( + model: str, custom_llm_provider: Optional[str] = None +) -> Optional[List[str]]: + """ + Get a list of supported regions for a given model and provider. + + Parameters: + model (str): The model name to be checked. + custom_llm_provider (Optional[str]): The provider to be checked. + """ + try: + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, custom_llm_provider=custom_llm_provider + ) + + model_info = _get_model_info_helper( + model=model, custom_llm_provider=custom_llm_provider + ) + + supported_regions = model_info.get("supported_regions", None) + if supported_regions is None: + return None + + ######################################################### + # Ensure only list supported regions are returned + ######################################################### + if isinstance(supported_regions, list): + return supported_regions + else: + return None + except Exception as e: + verbose_logger.debug( + f"Model not found or error in checking supported_regions support. You passed model={model}, custom_llm_provider={custom_llm_provider}. Error: {str(e)}" + ) + return None + + def supports_embedding_image_input( model: str, custom_llm_provider: Optional[str] = None ) -> bool: @@ -2836,10 +2873,10 @@ def pre_process_non_default_params( if "response_format" in non_default_params: if provider_config is not None: - non_default_params[ - "response_format" - ] = provider_config.get_json_schema_from_pydantic_object( - response_format=non_default_params["response_format"] + non_default_params["response_format"] = ( + provider_config.get_json_schema_from_pydantic_object( + response_format=non_default_params["response_format"] + ) ) else: non_default_params["response_format"] = type_to_response_format_param( @@ -2966,16 +3003,16 @@ def pre_process_optional_params( True # so that main.py adds the function call to the prompt ) if "tools" in non_default_params: - optional_params[ - "functions_unsupported_model" - ] = non_default_params.pop("tools") + optional_params["functions_unsupported_model"] = ( + non_default_params.pop("tools") + ) non_default_params.pop( "tool_choice", None ) # causes ollama requests to hang elif "functions" in non_default_params: - optional_params[ - "functions_unsupported_model" - ] = non_default_params.pop("functions") + optional_params["functions_unsupported_model"] = ( + non_default_params.pop("functions") + ) elif ( litellm.add_function_to_prompt ): # if user opts to add it to prompt instead @@ -4058,9 +4095,9 @@ def _count_characters(text: str) -> int: def get_response_string(response_obj: Union[ModelResponse, ModelResponseStream]) -> str: - _choices: Union[ - List[Union[Choices, StreamingChoices]], List[StreamingChoices] - ] = response_obj.choices + _choices: Union[List[Union[Choices, StreamingChoices]], List[StreamingChoices]] = ( + response_obj.choices + ) response_str = "" for choice in _choices: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 888097582d..3cda936572 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23,6 +23,13 @@ "search_context_size_medium": 0.0, "search_context_size_high": 0.0 }, + "supported_regions": [ + "global", + "us-west-2", + "eu-west-1", + "ap-southeast-1", + "ap-northeast-1" + ], "deprecation_date": "date when the model becomes deprecated in the format YYYY-MM-DD" }, "omni-moderation-latest": { @@ -6860,6 +6867,9 @@ "supported_output_modalities": [ "text" ], + "supported_regions": [ + "global" + ], "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", "supports_parallel_function_calling": true, "supports_web_search": true diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index 8ae49a2aa7..b43e5a113c 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -562,3 +562,68 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix): assert endpoint == expected_endpoint assert url == expected_url + + +@pytest.mark.parametrize( + "supported_regions, expected_result", + [ + (None, False), # get_supported_regions returns None + ([], False), # empty list, no global region + (["us-central1"], False), # only regional, no global + (["global"], True), # only global region + (["global", "us-central1"], True), # global and other regions + ( + ["us-central1", "global", "europe-west1"], + True, + ), # global among multiple regions + ], +) +def test_is_global_only_vertex_model(supported_regions, expected_result): + """Test is_global_only_vertex_model with various supported regions scenarios""" + from litellm.llms.vertex_ai.common_utils import is_global_only_vertex_model + + with patch("litellm.utils.get_supported_regions") as mock_get_supported_regions: + mock_get_supported_regions.return_value = supported_regions + + result = is_global_only_vertex_model("test-model") + + assert result == expected_result + mock_get_supported_regions.assert_called_once_with( + model="test-model", custom_llm_provider="vertex_ai" + ) + + +@pytest.mark.parametrize( + "model_is_global_only, vertex_region, expected_region", + [ + (True, None, "global"), # Global-only model with no region specified + (True, "us-central1", "global"), # Global-only model overrides specified region + (True, "europe-west1", "global"), # Global-only model overrides any region + (False, None, "us-central1"), # Non-global model defaults to us-central1 + ( + False, + "europe-west1", + "europe-west1", + ), # Non-global model uses specified region + (False, "us-east1", "us-east1"), # Non-global model uses specified region + ], +) +def test_get_vertex_region_global_only_model( + model_is_global_only, vertex_region, expected_region +): + """Test get_vertex_region ensures global-only models default to 'global' region""" + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + vertex_base = VertexBase() + + with patch( + "litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model" + ) as mock_is_global_only: + mock_is_global_only.return_value = model_is_global_only + + result = vertex_base.get_vertex_region( + vertex_region=vertex_region, model="test-model" + ) + + assert result == expected_region + mock_is_global_only.assert_called_once_with("test-model") diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index eccd11845a..156840d669 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -451,6 +451,12 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): ], }, }, + "supported_regions": { + "type": "array", + "items": { + "type": "string", + }, + }, "search_context_cost_per_query": { "type": "object", "properties": {