Merge pull request #26076 from BerriAI/litellm_vertex_model_garden_xai_openapi

feat(vertex_ai): Model Garden OpenAPI for publisher model ids
2026-10-07 02:59:05 +00:00 · 2026-05-01 10:15:58 -07:00 · 2026-05-01 10:15:58 -07:00 · a155ea1e8a
commit a155ea1e8a
parent 32ab390e7e e656b2a47b
5 changed files with 203 additions and 8 deletions
--- a/litellm/llms/vertex_ai/common_utils.py
+++ b/litellm/llms/vertex_ai/common_utils.py
@ -97,7 +97,7 @@ def get_vertex_ai_model_route(
    Determine which handler to use for a Vertex AI model based on the model name.

    Args:
-        model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "openai/gpt-oss-120b")
+        model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "xai/grok-4.1-fast-non-reasoning")
        litellm_params: Optional litellm parameters dict that may contain base_model for routing

    Returns:
@ -113,7 +113,7 @@ def get_vertex_ai_model_route(
        >>> get_vertex_ai_model_route("gemma/gemma-3-12b-it")
        VertexAIModelRoute.GEMMA

-        >>> get_vertex_ai_model_route("openai/gpt-oss-120b")
+        >>> get_vertex_ai_model_route("xai/grok-4.1-fast-non-reasoning")
        VertexAIModelRoute.MODEL_GARDEN

        >>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"})
@ -149,8 +149,11 @@ def get_vertex_ai_model_route(
    if "gemma/" in model:
        return VertexAIModelRoute.GEMMA

-    # Check for model garden openai models
-    if "openai" in model:
+    # Check for model garden OpenAI-compatible publisher models.
+    # Examples:
+    # - openai/gpt-oss-120b-maas
+    # - xai/grok-4.1-fast-non-reasoning
+    if "openai" in model or model.startswith("xai/"):
        return VertexAIModelRoute.MODEL_GARDEN

    # Check for gemini models
@ -256,8 +259,8 @@ def get_vertex_base_model_name(model: str) -> str:
        >>> get_vertex_base_model_name("gemma/gemma-3-12b-it")
        "gemma-3-12b-it"

-        >>> get_vertex_base_model_name("openai/gpt-oss-120b")
-        "gpt-oss-120b"
+        >>> get_vertex_base_model_name("xai/grok-4.1-fast-non-reasoning")
+        "grok-4.1-fast-non-reasoning"

        >>> get_vertex_base_model_name("1234567890")
        "1234567890"
--- a/litellm/llms/vertex_ai/vertex_model_garden/main.py
+++ b/litellm/llms/vertex_ai/vertex_model_garden/main.py
@ -27,6 +27,17 @@ from ..common_utils import VertexAIError, get_vertex_base_model_name
 from ..vertex_llm_base import VertexBase


+def _vertex_model_garden_model_id_in_json_body(model: str) -> bool:
+    """
+    Vertex catalog / publisher models are addressed as publisher/model (e.g.
+    xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body.
+
+    Deployed Model Garden endpoints are typically a single segment (often numeric)
+    and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field.
+    """
+    return "/" in model
+
+
 def create_vertex_url(
    vertex_location: str,
    vertex_project: str,
@ -34,8 +45,13 @@ def create_vertex_url(
    model: str,
    api_base: Optional[str] = None,
 ) -> str:
-    """Return the base url for the vertex garden models"""
+    """Return the api base for vertex model garden (without /chat/completions)."""
    base_url = get_vertex_base_url(vertex_location)
+    if _vertex_model_garden_model_id_in_json_body(model):
+        return (
+            f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}"
+            "/endpoints/openapi"
+        )
    return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}"


@ -129,7 +145,10 @@ class VertexAIModelGardenModels(VertexBase):
                vertex_location=vertex_location or "us-central1",
                vertex_api_version="v1beta1",
            )
-            model = ""
+            # Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route).
+            # Single-segment endpoint ids: model is encoded in the URL path; body model stays empty.
+            if not _vertex_model_garden_model_id_in_json_body(model):
+                model = ""
            return openai_like_chat_completions.completion(
                model=model,
                messages=messages,
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@ -33337,6 +33337,72 @@
        "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
        "supports_reasoning": true
    },
+    "vertex_ai/xai/grok-4.1-fast-non-reasoning": {
+        "cache_read_input_token_cost": 5e-08,
+        "input_cost_per_token": 2e-07,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 5e-07,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.1-fast-reasoning": {
+        "cache_read_input_token_cost": 5e-08,
+        "input_cost_per_token": 2e-07,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 5e-07,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.20-non-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.20-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
    "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
        "input_cost_per_token": 2.5e-07,
        "litellm_provider": "vertex_ai-qwen_models",
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@ -33391,6 +33391,72 @@
        "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
        "supports_reasoning": true
    },
+    "vertex_ai/xai/grok-4.1-fast-non-reasoning": {
+        "cache_read_input_token_cost": 5e-08,
+        "input_cost_per_token": 2e-07,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 5e-07,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.1-fast-reasoning": {
+        "cache_read_input_token_cost": 5e-08,
+        "input_cost_per_token": 2e-07,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 5e-07,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.20-non-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.20-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
    "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
        "input_cost_per_token": 2.5e-07,
        "litellm_provider": "vertex_ai-qwen_models",
--- a/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py
+++ b/tests/test_litellm/llms/vertex_ai/test_vertex_model_garden_openapi.py
@ -0,0 +1,41 @@
+"""Vertex Model Garden: OpenAPI base URL for publisher/model ids vs per-endpoint path."""
+
+import pytest
+
+from litellm.llms.vertex_ai.vertex_model_garden.main import (
+    _vertex_model_garden_model_id_in_json_body,
+    create_vertex_url,
+)
+
+
+@pytest.mark.parametrize(
+    "model,expect_openapi_base",
+    [
+        ("xai/grok-4.1-fast-reasoning", True),
+        ("openai/foo/bar", True),
+        ("5464397967697903616", False),
+        ("gpt-oss-20b-maas", False),
+    ],
+)
+def test_create_vertex_url_openapi_vs_deployed_endpoint(
+    model: str, expect_openapi_base: bool
+) -> None:
+    url = create_vertex_url(
+        vertex_location="us-central1",
+        vertex_project="my-project",
+        stream=False,
+        model=model,
+    )
+    if expect_openapi_base:
+        assert "/v1/projects/my-project/locations/us-central1/endpoints/openapi" in url
+    else:
+        assert (
+            "/v1beta1/projects/my-project/locations/us-central1/endpoints/"
+            f"{model}" in url
+        )
+        assert "openapi" not in url
+
+
+def test_model_id_in_json_body_heuristic() -> None:
+    assert _vertex_model_garden_model_id_in_json_body("xai/grok-4.1-fast-reasoning") is True
+    assert _vertex_model_garden_model_id_in_json_body("5464397967697903616") is False