# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
from typing import Any, Optional
import pydantic
from typing_extensions import override
from oumi.core.configs import GenerationParams, ModelParams, RemoteParams
from oumi.core.configs.params.guided_decoding_params import GuidedDecodingParams
from oumi.core.types.conversation import Conversation
from oumi.inference.remote_inference_engine import RemoteInferenceEngine
[docs]
class GoogleVertexInferenceEngine(RemoteInferenceEngine):
"""Engine for running inference against Google Vertex AI."""
_API_URL_TEMPLATE = (
"https://{region}-aiplatform.googleapis.com/v1beta1/projects/"
"{project_id}/locations/{region}/endpoints/openapi/chat/completions"
)
"""The API URL template for the GCP project. Used when no `api_url` is provided."""
_DEFAULT_PROJECT_ID_ENV_KEY: str = "PROJECT_ID"
"""The default project ID environment key for the GCP project."""
_DEFAULT_REGION_ENV_KEY: str = "REGION"
"""The default region environment key for the GCP project."""
_project_id: Optional[str] = None
"""The project ID for the GCP project."""
_region: Optional[str] = None
"""The region for the GCP project."""
def __init__(
self,
model_params: ModelParams,
*,
generation_params: Optional[GenerationParams] = None,
remote_params: Optional[RemoteParams] = None,
project_id_env_key: Optional[str] = None,
region_env_key: Optional[str] = None,
project_id: Optional[str] = None,
region: Optional[str] = None,
):
"""Initializes the inference Engine.
Args:
model_params: The model parameters to use for inference.
generation_params: The generation parameters to use for inference.
remote_params: The remote parameters to use for inference.
project_id_env_key: The environment variable key name for the project ID.
region_env_key: The environment variable key name for the region.
project_id: The project ID to use for inference.
region: The region to use for inference.
"""
super().__init__(
model_params=model_params,
generation_params=generation_params,
remote_params=remote_params,
)
if project_id and project_id_env_key:
raise ValueError(
"You cannot set both `project_id` and `project_id_env_key`."
)
if region and region_env_key:
raise ValueError("You cannot set both `region` and `region_env_key`.")
self._project_id_env_key = (
project_id_env_key or self._DEFAULT_PROJECT_ID_ENV_KEY
)
self._region_env_key = region_env_key or self._DEFAULT_REGION_ENV_KEY
self._project_id = project_id
self._region = region
@override
def _set_required_fields_for_inference(self, remote_params: RemoteParams) -> None:
"""Set required fields for inference."""
if (
not remote_params.api_url
and not self._remote_params.api_url
and not self.base_url
):
if self._project_id and self._region:
project_id = self._project_id
region = self._region
elif os.getenv(self._project_id_env_key) and os.getenv(
self._region_env_key
):
project_id = os.getenv(self._project_id_env_key)
region = os.getenv(self._region_env_key)
else:
raise ValueError(
"This inference engine requires that either `api_url` is set in "
"`RemoteParams` or that both `project_id` and `region` are set. "
"You can set the `project_id` and `region` when "
"constructing a GoogleVertexInferenceEngine, "
f"or as environment variables: `{self._project_id_env_key}` and "
f"`{self._region_env_key}`."
)
remote_params.api_url = self._API_URL_TEMPLATE.format(
project_id=project_id,
region=region,
)
super()._set_required_fields_for_inference(remote_params)
@override
def _get_api_key(self, remote_params: RemoteParams) -> str:
"""Gets the authentication token for GCP."""
try:
from google.auth import default # pyright: ignore[reportMissingImports]
from google.auth.transport.requests import ( # pyright: ignore[reportMissingImports]
Request,
)
from google.oauth2 import ( # pyright: ignore[reportMissingImports]
service_account,
)
except ModuleNotFoundError:
raise RuntimeError(
"Google-auth is not installed. "
"Please install oumi with GCP extra:`pip install oumi[gcp]`, "
"or install google-auth with `pip install google-auth`."
)
if remote_params.api_key:
credentials = service_account.Credentials.from_service_account_file(
filename=remote_params.api_key,
scopes=["https://www.googleapis.com/auth/cloud-platform"],
)
else:
credentials, _ = default(
scopes=["https://www.googleapis.com/auth/cloud-platform"]
)
credentials.refresh(Request()) # type: ignore
return credentials.token # type: ignore
@override
def _get_request_headers(
self, remote_params: Optional[RemoteParams]
) -> dict[str, str]:
"""Gets the request headers for GCP."""
if not remote_params:
raise ValueError("Remote params are required for GCP inference.")
headers = {
"Authorization": f"Bearer {self._get_api_key(remote_params)}",
"Content-Type": "application/json",
}
return headers
@override
def _default_remote_params(self) -> RemoteParams:
"""Returns the default remote parameters."""
return RemoteParams(num_workers=10, politeness_policy=60.0)
@override
def _convert_conversation_to_api_input(
self,
conversation: Conversation,
generation_params: GenerationParams,
model_params: ModelParams,
) -> dict[str, Any]:
"""Converts a conversation to an OpenAI input.
Documentation: https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/call-vertex-using-openai-library
Args:
conversation: The conversation to convert.
generation_params: Parameters for generation during inference.
model_params: Model parameters to use during inference.
Returns:
Dict[str, Any]: A dictionary representing the Vertex input.
"""
api_input = {
"model": model_params.model_name,
"messages": self._get_list_of_message_json_dicts(
conversation.messages, group_adjacent_same_role_turns=True
),
"max_completion_tokens": generation_params.max_new_tokens,
"temperature": generation_params.temperature,
"top_p": generation_params.top_p,
"n": 1, # Number of completions to generate for each prompt.
"seed": generation_params.seed,
"logit_bias": generation_params.logit_bias,
}
if generation_params.stop_strings:
api_input["stop"] = generation_params.stop_strings
if generation_params.guided_decoding:
api_input["response_format"] = _convert_guided_decoding_config_to_api_input(
generation_params.guided_decoding
)
return api_input
[docs]
@override
def get_supported_params(self) -> set[str]:
"""Returns a set of supported generation parameters for this engine."""
return {
"guided_decoding",
"logit_bias",
"max_new_tokens",
"seed",
"stop_strings",
"temperature",
"top_p",
}
#
# Helper functions
#
def _convert_guided_decoding_config_to_api_input(
guided_config: GuidedDecodingParams,
) -> dict:
"""Converts a guided decoding configuration to an API input."""
if guided_config.json is None:
raise ValueError(
"Only JSON schema guided decoding is supported, got '%s'",
guided_config,
)
json_schema = guided_config.json
if isinstance(json_schema, type) and issubclass(json_schema, pydantic.BaseModel):
schema_name = json_schema.__name__
schema_value = json_schema.model_json_schema()
elif isinstance(json_schema, dict):
# Use a generic name if no schema is provided.
schema_name = "Response"
schema_value = json_schema
elif isinstance(json_schema, str):
# Use a generic name if no schema is provided.
schema_name = "Response"
# Try to parse as JSON string
schema_value = json.loads(json_schema)
else:
raise ValueError(
f"Got unsupported JSON schema type: {type(json_schema)}"
"Please provide a Pydantic model or a JSON schema as a "
"string or dict."
)
return {
"type": "json_schema",
"json_schema": {
"name": schema_name,
"schema": _replace_refs_in_schema(schema_value),
},
}
def _replace_refs_in_schema(schema: dict) -> dict:
"""Replace $ref references in a JSON schema with their actual definitions.
Args:
schema: The JSON schema dictionary
Returns:
dict: Schema with all references replaced by their definitions and $defs removed
"""
def _get_ref_value(ref: str) -> dict:
# Remove the '#/' prefix and split into parts
parts = ref.replace("#/", "").split("/")
# Navigate through the schema to get the referenced value
current = schema
for part in parts:
current = current[part]
return current.copy() # Return a copy to avoid modifying the original
def _replace_refs(obj: dict) -> dict:
if not isinstance(obj, dict):
return obj
result = {}
for key, value in obj.items():
if key == "$ref":
# If we find a $ref, replace it with the actual value
return _replace_refs(_get_ref_value(value))
elif isinstance(value, dict):
result[key] = _replace_refs(value)
elif isinstance(value, list):
result[key] = [
_replace_refs(item) if isinstance(item, dict) else item
for item in value
]
else:
result[key] = value
return result
# Replace all references first
resolved = _replace_refs(schema.copy())
# Remove the $defs key if it exists
if "$defs" in resolved:
del resolved["$defs"]
return resolved
