tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 16:24:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #19104 from Chesars/fix/vertex-ai-zai-org-global-region

Browse Source 
feat(vertex_ai): route region for partner models and add GLM support
This commit is contained in:
Cesar Garcia
2026-03-11 15:19:33 -03:00
committed by GitHub
parent 0c95d415e1 3948513a4c
commit cb24b8b05e
7 changed files with 163 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/providers/vertex_partner.md
+74
View File
@@ -11,6 +11,7 @@ import TabItem from '@theme/TabItem';
|----------|---------------|---------------|
| Anthropic (Claude) | `vertex_ai/claude-*` | [Vertex AI - Anthropic Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude) |
| DeepSeek | `vertex_ai/deepseek-ai/{MODEL}` | [Vertex AI - DeepSeek Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/deepseek) |
| ZAI (GLM) | `vertex_ai/zai-org/{MODEL}` | [Vertex AI - GLM Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/zaiorg/glm-47) |
| Meta/Llama | `vertex_ai/meta/{MODEL}` | [Vertex AI - Meta Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) |
| Mistral | `vertex_ai/mistral-*` | [Vertex AI - Mistral Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral) |
| AI21 (Jamba) | `vertex_ai/jamba-*` | [Vertex AI - AI21 Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/ai21) |
@@ -226,6 +227,79 @@ ModelResponse(
|------------------|------------------------------|
| vertex_ai/deepseek-ai/deepseek-r1-0528-maas | `completion('vertex_ai/deepseek-ai/deepseek-r1-0528-maas', messages)` |
## VertexAI ZAI (GLM)
| Property | Details |
|----------|---------|
| Provider Route | `vertex_ai/zai-org/{MODEL}` |
| Vertex Documentation | [Vertex AI - GLM Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/zaiorg/glm-47) |
**LiteLLM Supports all Vertex AI GLM Models.** Ensure you use the `vertex_ai/zai-org/` prefix for all Vertex AI GLM models.
| Model Name | Usage |
|------------|-------|
| vertex_ai/zai-org/glm-4.7-maas | `completion('vertex_ai/zai-org/glm-4.7-maas', messages)` |
#### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ""
response = completion(
model="vertex_ai/zai-org/glm-4.7-maas",
messages=[{"role": "user", "content": "hi"}],
vertex_project="your-vertex-project",
# vertex_location routes to "global"
)
print("\nModel Response", response)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: glm-4.7
litellm_params:
model: vertex_ai/zai-org/glm-4.7-maas
vertex_project: "my-project"
# vertex_location routes to "global"
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "glm-4.7",
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
## VertexAI Meta/Llama API
litellm/llms/vertex_ai/common_utils.py
-19
View File
@@ -940,25 +940,6 @@ def construct_target_url(
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
class VertexAIModelInfo(BaseLLMModelInfo):
def get_token_counter(self) -> Optional[BaseTokenCounter]:
litellm/llms/vertex_ai/vertex_llm_base.py
+21 -3
View File
@@ -21,7 +21,6 @@ from .common_utils import (
all_gemini_url_modes,
get_vertex_base_model_name,
get_vertex_base_url,
is_global_only_vertex_model,
)
GOOGLE_IMPORT_ERROR_MESSAGE = (
@@ -49,8 +48,27 @@ class VertexBase:
self.async_handler: Optional[AsyncHTTPHandler] = None
def get_vertex_region(self, vertex_region: Optional[str], model: str) -> str:
if is_global_only_vertex_model(model):
return "global"
import litellm
# Try to get supported_regions directly from model_cost
# Check both with and without vertex_ai/ prefix
model_key = f"vertex_ai/{model}" if not model.startswith("vertex_ai/") else model
model_info = litellm.model_cost.get(model_key, {})
supported_regions = model_info.get("supported_regions")
if supported_regions and len(supported_regions) > 0:
# If user didn't specify region, use the first supported region
if vertex_region is None:
return supported_regions[0]
# If user specified a region not supported by this model, override it
if vertex_region not in supported_regions:
verbose_logger.warning(
"Vertex AI model '%s' does not support region '%s' "
"(supported: %s). Routing to '%s'.",
model, vertex_region, supported_regions, supported_regions[0],
)
return supported_regions[0]
return vertex_region
return vertex_region or "us-central1"
def load_auth(
litellm/model_prices_and_context_window_backup.json
+4
View File
@@ -33840,6 +33840,9 @@
"mode": "chat",
"output_cost_per_token": 2.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
@@ -33854,6 +33857,7 @@
"mode": "chat",
"output_cost_per_token": 3.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#glm-models",
"supported_regions": ["global"],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
model_prices_and_context_window.json
+4
View File
@@ -33840,6 +33840,9 @@
"mode": "chat",
"output_cost_per_token": 2.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#partner-models",
"supported_regions": [
"global"
],
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
@@ -33854,6 +33857,7 @@
"mode": "chat",
"output_cost_per_token": 3.2e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#glm-models",
"supported_regions": ["global"],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
+27 -47
View File
@@ -558,69 +558,49 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
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",
"model_cost_entry, 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
# Model with supported_regions=["global"], no user region -> use "global"
({"supported_regions": ["global"]}, None, "global"),
# Model with supported_regions=["global"], user passes unsupported region -> override to "global"
({"supported_regions": ["global"]}, "us-central1", "global"),
# Model with supported_regions=["global"], user passes unsupported region -> override to "global"
({"supported_regions": ["global"]}, "europe-west1", "global"),
# Model with supported_regions=["us-west2"], no user region -> use "us-west2"
({"supported_regions": ["us-west2"]}, None, "us-west2"),
# Model with supported_regions=["us-west2", "us-central1"], user passes supported region -> respect it
({"supported_regions": ["us-west2", "us-central1"]}, "us-central1", "us-central1"),
# Model with supported_regions=["us-west2", "us-central1"], user passes unsupported region -> override
({"supported_regions": ["us-west2", "us-central1"]}, "europe-west1", "us-west2"),
# No model_cost entry, no user region -> default us-central1
({}, None, "us-central1"),
# No model_cost entry, user specifies region -> use specified region
({}, "europe-west1", "europe-west1"),
# No model_cost entry, user specifies region -> use specified region
({}, "us-east1", "us-east1"),
],
)
def test_get_vertex_region_global_only_model(
model_is_global_only, vertex_region, expected_region
model_cost_entry, vertex_region, expected_region
):
"""Test get_vertex_region ensures global-only models default to 'global' region"""
"""Test get_vertex_region resolves region from model_cost supported_regions"""
import litellm
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
with patch.dict(
litellm.model_cost,
{"vertex_ai/test-model": model_cost_entry},
clear=False,
):
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")
def test_vertex_filter_format_uri():
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/qwen/test_vertex_ai_qwen_global_endpoint.py
+33 -70
View File
@@ -2,8 +2,8 @@
Tests for Vertex AI Qwen MaaS models that require the global endpoint.
These tests verify that:
1. Qwen models are correctly identified as global-only models
2. The correct global URL is constructed (https://aiplatform.googleapis.com)
1. The correct global URL is constructed (https://aiplatform.googleapis.com)
2. The get_vertex_region method resolves regions from model_cost supported_regions
3. The completion() and responses() API work with Qwen models
"""
@@ -19,7 +19,6 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm
from litellm.llms.vertex_ai.common_utils import is_global_only_vertex_model
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VertexPartnerProvider
@@ -48,66 +47,36 @@ def clean_vertex_env():
os.environ[var] = value
class TestQwenGlobalOnlyDetection:
"""Test that Qwen models are correctly identified as global-only."""
@pytest.mark.parametrize(
"model",
[
"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
"vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas",
"vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas",
"vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas",
],
)
def test_qwen_models_are_global_only(self, model):
"""Test that Qwen MaaS models are identified as global-only."""
# This test requires the model_cost to have supported_regions: ["global"]
# If the model is not in model_cost, it should return False (fallback behavior)
result = is_global_only_vertex_model(model)
# Note: This will return True only if the model is in model_cost with supported_regions: ["global"]
# If running without the updated model_cost, this may return False
assert isinstance(result, bool)
def test_non_global_model_returns_false(self):
"""Test that non-global models return False."""
result = is_global_only_vertex_model("vertex_ai/gemini-1.5-pro")
assert result is False
def test_unknown_model_returns_false(self):
"""Test that unknown models return False (fallback behavior)."""
result = is_global_only_vertex_model("vertex_ai/unknown-model-xyz")
assert result is False
class TestVertexBaseGetVertexRegion:
"""Test the get_vertex_region method."""
"""Test the get_vertex_region method using model_cost lookup."""
def test_global_only_model_returns_global(self):
"""Test that global-only models return 'global' regardless of input."""
def test_global_model_no_user_region_returns_global(self):
"""Test that global-only models return 'global' when user doesn't specify region."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=True,
):
result = vertex_base.get_vertex_region(
vertex_region="us-central1",
model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
)
assert result == "global"
def test_global_only_model_with_none_returns_global(self):
"""Test that global-only models return 'global' even with None input."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=True,
with patch.dict(
litellm.model_cost,
{"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region=None,
model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
model="qwen/qwen3-next-80b-a3b-instruct-maas",
)
assert result == "global"
def test_global_model_with_unsupported_user_region_overrides(self):
"""Test that unsupported user region is overridden for global-only models."""
vertex_base = VertexBase()
with patch.dict(
litellm.model_cost,
{"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}},
clear=False,
):
result = vertex_base.get_vertex_region(
vertex_region="us-central1",
model="qwen/qwen3-next-80b-a3b-instruct-maas",
)
assert result == "global"
@@ -115,13 +84,10 @@ class TestVertexBaseGetVertexRegion:
"""Test that non-global models use the provided region."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=False,
):
with patch.dict(litellm.model_cost, {}, clear=False):
result = vertex_base.get_vertex_region(
vertex_region="europe-west1",
model="vertex_ai/gemini-1.5-pro",
model="gemini-1.5-pro",
)
assert result == "europe-west1"
@@ -129,13 +95,10 @@ class TestVertexBaseGetVertexRegion:
"""Test that non-global models with None region fallback to us-central1."""
vertex_base = VertexBase()
with patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=False,
):
with patch.dict(litellm.model_cost, {}, clear=False):
result = vertex_base.get_vertex_region(
vertex_region=None,
model="vertex_ai/gemini-1.5-pro",
model="unknown-model-xyz",
)
assert result == "us-central1"
@@ -217,15 +180,15 @@ async def test_vertex_ai_qwen_global_endpoint_url():
client, "post", side_effect=mock_post_func
) as mock_post, patch.object(
VertexLLM, "_ensure_access_token", return_value=("fake-token", "test-project")
), patch(
"litellm.llms.vertex_ai.vertex_llm_base.is_global_only_vertex_model",
return_value=True,
), patch.dict(
litellm.model_cost,
{"vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": {"supported_regions": ["global"]}},
clear=False,
):
response = await litellm.acompletion(
model="vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas",
messages=[{"role": "user", "content": "Hello"}],
vertex_ai_project="test-project",
vertex_ai_location="us-central1",
client=client,
)