"""
API Handler for calling Vertex AI Gemma Models
These models use a custom prediction endpoint format that wraps messages in 'instances'
with @requestFormat: "chatCompletions" and returns responses wrapped in 'predictions'.
Usage:
response = litellm.completion(
model="vertex_ai/gemma/gemma-3-12b-it-1222199011122",
messages=[{"role": "user", "content": "What is machine learning?"}],
vertex_project="your-project-id",
vertex_location="us-central1",
)
Sent to this route when `model` is in the format `vertex_ai/gemma/{MODEL_NAME}`
The API expects a custom endpoint URL format:
https://{ENDPOINT_NUMBER}.{location}-{REGION_NUMBER}.prediction.vertexai.goog/v1/projects/{PROJECT_ID}/locations/{location}/endpoints/{ENDPOINT_ID}:predict
"""
from typing import Callable, Optional, Union
import httpx # type: ignore
from litellm.utils import ModelResponse
from ..common_utils import VertexAIError, get_vertex_base_model_name
from ..vertex_llm_base import VertexBase
class VertexAIGemmaModels(VertexBase):
def __init__(self) -> None:
pass
def completion(
self,
model: str,
messages: list,
model_response: ModelResponse,
print_verbose: Callable,
encoding,
logging_obj,
api_base: Optional[str],
optional_params: dict,
custom_prompt_dict: dict,
headers: Optional[dict],
timeout: Union[float, httpx.Timeout],
litellm_params: dict,
vertex_project=None,
vertex_location=None,
vertex_credentials=None,
logger_fn=None,
acompletion: bool = False,
client=None,
):
"""
Handles calling Vertex AI Gemma Models
Sent to this route when `model` is in the format `vertex_ai/gemma/{MODEL_NAME}`
"""
try:
import vertexai
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexLLM,
)
from litellm.llms.vertex_ai.vertex_gemma_models.transformation import (
VertexGemmaConfig,
)
except Exception as e:
raise VertexAIError(
status_code=400,
message=f"""vertexai import failed please run `pip install -U "google-cloud-aiplatform>=1.38"`. Got error: {e}""",
)
if not (
hasattr(vertexai, "preview") or hasattr(vertexai.preview, "language_models")
):
raise VertexAIError(
status_code=400,
message="""Upgrade vertex ai. Run `pip install "google-cloud-aiplatform>=1.38"`""",
)
try:
model = get_vertex_base_model_name(model=model)
vertex_httpx_logic = VertexLLM()
access_token, project_id = vertex_httpx_logic._ensure_access_token(
credentials=vertex_credentials,
project_id=vertex_project,
custom_llm_provider="vertex_ai",
)
gemma_transformation = VertexGemmaConfig()
## CONSTRUCT API BASE
stream: bool = optional_params.get("stream", False) or False
optional_params["stream"] = stream
# If api_base is not provided, it should be set as an environment variable
# or passed explicitly because the endpoint URL is unique per deployment
if api_base is None:
raise VertexAIError(
status_code=400,
message="api_base is required for Vertex AI Gemma models. Please provide the full endpoint URL.",
)
# Check if we need to append :predict
if not api_base.endswith(":predict"):
_, api_base = self._check_custom_proxy(
api_base=api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=stream,
auth_header=None,
url=api_base,
)
# If api_base already ends with :predict, use it as-is
# Use the custom transformation handler for gemma models
return gemma_transformation.completion(
model=model,
messages=messages,
api_base=api_base,
api_key=access_token,
custom_prompt_dict=custom_prompt_dict,
model_response=model_response,
print_verbose=print_verbose,
logging_obj=logging_obj,
optional_params=optional_params,
acompletion=acompletion,
litellm_params=litellm_params,
logger_fn=logger_fn,
client=client,
timeout=timeout,
encoding=encoding,
custom_llm_provider="vertex_ai",
)
except Exception as e:
if hasattr(e, "status_code"):
raise e
raise VertexAIError(status_code=500, message=str(e))