ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix: Adds support for choosing the default region based on where the model is available (#11566)

Browse Source 
* fix: vtx default region for global only models

* track gemini-2.5-pro-preview-05-06

* fix is_global_only_vertex_model

* test_is_global_only_vertex_model

* test_get_vertex_region_global_only_model

* fix json format

* fix get_supported_regions
This commit is contained in:
Ishaan Jaff 2025-06-09 18:29:44 -07:00 committed by GitHub
parent 230dd70604
commit 9241fca2f5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 180 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
litellm/llms/vertex_ai/common_utils.py
View File

@ -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

16
litellm/llms/vertex_ai/vertex_llm_base.py
View File

@ -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"] = (

10
litellm/model_prices_and_context_window_backup.json
View File

@ -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

75
litellm/utils.py
View File

@ -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:

10
model_prices_and_context_window.json
View File

@ -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

65
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
View File

@ -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")

6
tests/test_litellm/test_utils.py
View File

@ -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": {