feat: GenAI Client(evals) - Add support for Vertex Model Garden MaaS models.

PiperOrigin-RevId: 806024875

Author: jsondai

tests/unit/vertexai/genai/test_evals.py

@@ -1102,7 +1102,7 @@ def test_run_inference_with_litellm_import_error(self, mock_api_client_fixture):
         prompt_df = pd.DataFrame([{"prompt": "test"}])
         with pytest.raises(
             ImportError,
-            match="The 'litellm' library is required to use third-party models",
+            match="The 'litellm' library is required to use this model.",
         ):
             evals_module.run_inference(model="gpt-4o", src=prompt_df)
 

vertexai/_genai/_evals_common.py

@@ -29,6 +29,7 @@
 import pandas as pd
 from tqdm import tqdm
 
+from . import _evals_constant
 from . import _evals_data_converters
 from . import _evals_metric_handlers
 from . import _evals_utils
@@ -370,6 +371,14 @@ def _run_litellm_inference(
     return responses
 
 
+def _is_litellm_vertex_maas_model(model: str) -> bool:
+    """Checks if the model is a Vertex MAAS model to be handled by LiteLLM."""
+    return any(
+        model.startswith(prefix)
+        for prefix in _evals_constant.SUPPORTED_VERTEX_MAAS_MODEL_PREFIXES
+    )
+
+
 def _is_litellm_model(model: str) -> bool:
     """Checks if the model name corresponds to a valid LiteLLM model name."""
     return model in litellm.utils.get_valid_models(model)
@@ -431,47 +440,64 @@ def _run_inference_internal(
                 }
                 processed_responses.append(json.dumps(error_payload))
         responses = processed_responses
+    elif callable(model):
+        logger.info("Running inference with custom callable function.")
+        custom_responses_raw = _run_custom_inference(
+            model_fn=model, prompt_dataset=prompt_dataset
+        )
+        processed_custom_responses = []
+        for resp_item in custom_responses_raw:
+            if isinstance(resp_item, str):
+                processed_custom_responses.append(resp_item)
+            elif isinstance(resp_item, dict) and "error" in resp_item:
+                processed_custom_responses.append(json.dumps(resp_item))
+            else:
+                try:
+                    processed_custom_responses.append(json.dumps(resp_item))
+                except TypeError:
+                    processed_custom_responses.append(str(resp_item))
+        responses = processed_custom_responses
     elif isinstance(model, str):
         if litellm is None:
             raise ImportError(
-                "The 'litellm' library is required to use third-party models."
+                "The 'litellm' library is required to use this model."
                 " Please install it using 'pip install"
                 " google-cloud-aiplatform[evaluation]'."
             )
-        if _is_litellm_model(model):
-            logger.info("Running inference with LiteLLM for model: %s", model)
-            raw_responses = _run_litellm_inference(  # type: ignore[assignment]
-                model=model, prompt_dataset=prompt_dataset
+
+        processed_model_id = model
+        if model.startswith("vertex_ai/"):
+            # Already correctly prefixed for LiteLLM's Vertex AI provider
+            pass
+        elif _is_litellm_vertex_maas_model(model):
+            processed_model_id = f"vertex_ai/{model}"
+            logger.info(
+                "Detected Vertex AI Model Garden managed MaaS model. "
+                "Using LiteLLM ID: %s",
+                processed_model_id,
             )
-            responses = [json.dumps(resp) for resp in raw_responses]
+        elif _is_litellm_model(model):
+            # Other LiteLLM supported model
+            logger.info("Running inference with LiteLLM for model: %s", model)
         else:
+            # Unsupported model string
             raise TypeError(
                 f"Unsupported string model name: {model}. Expecting a Gemini model"
-                " name (e.g., 'gemini-2.5-pro', 'projects/.../models/...') or a"
+                " name (e.g., 'gemini-1.5-pro', 'projects/.../models/...') or a"
                 " LiteLLM supported model name (e.g., 'openai/gpt-4o')."
                 " If using a third-party model via LiteLLM, ensure the"
                 " necessary environment variables are set (e.g., for OpenAI:"
                 " `os.environ['OPENAI_API_KEY'] = 'Your API Key'`). See"
                 " LiteLLM documentation for details:"
                 " https://docs.litellm.ai/docs/set_keys#environment-variables"
             )
-    elif callable(model):
-        logger.info("Running inference with custom callable function.")
-        custom_responses_raw = _run_custom_inference(
-            model_fn=model, prompt_dataset=prompt_dataset
+
+        logger.info("Running inference via LiteLLM for model: %s", processed_model_id)
+        raw_responses = _run_litellm_inference(
+            model=processed_model_id, prompt_dataset=prompt_dataset
         )
-        processed_custom_responses = []
-        for resp_item in custom_responses_raw:
-            if isinstance(resp_item, str):
-                processed_custom_responses.append(resp_item)
-            elif isinstance(resp_item, dict) and "error" in resp_item:
-                processed_custom_responses.append(json.dumps(resp_item))
-            else:
-                try:
-                    processed_custom_responses.append(json.dumps(resp_item))
-                except TypeError:
-                    processed_custom_responses.append(str(resp_item))
-        responses = processed_custom_responses
+        responses = [json.dumps(resp) for resp in raw_responses]
+
     else:
         raise TypeError(
             f"Unsupported model type: {type(model)}. Expecting string (model"

vertexai/_genai/_evals_constant.py

@@ -27,3 +27,15 @@
         "final_response_reference_free_v1",
     }
 )
+
+SUPPORTED_VERTEX_MAAS_MODEL_PREFIXES = frozenset(
+    {
+        "meta/",  # Meta/Llama
+        "deepseek-ai/",  # DeepSeek AI
+        "qwen/",  # Qwen
+        "openai/",  # OpenAI (GPT-OSS)
+        "claude-",  # Anthropic (Claude)
+        "mistral-",  # Mistral AI
+        "jamba-",  # AI21 (Jamba)
+    }
+)