diff --git a/agentplatform/_genai/client.py b/agentplatform/_genai/client.py index fb9238132e..449afd6e29 100644 --- a/agentplatform/_genai/client.py +++ b/agentplatform/_genai/client.py @@ -62,6 +62,10 @@ from agentplatform._genai import ( memory_banks as memory_banks_module, ) + from agentplatform._genai import ( + serving_profiles as serving_profiles_module, + ) + _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules" @@ -112,6 +116,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._sessions: Optional[ModuleType] = None self._sandboxes: Optional[ModuleType] = None self._memory_banks: Optional[ModuleType] = None + self._serving_profiles: Optional[ModuleType] = None @property @_common.experimental_warning( @@ -287,6 +292,17 @@ def memory_banks(self) -> "memory_banks_module.AsyncMemoryBanks": self._memory_banks = importlib.import_module(".memory_banks", __package__) return self._memory_banks.AsyncMemoryBanks(self._api_client) # type: ignore[no-any-return] + @property + def serving_profiles( + self, + ) -> "serving_profiles_module.AsyncServingProfiles": + if self._serving_profiles is None: + self._serving_profiles = importlib.import_module( + ".serving_profiles", + __package__, + ) + return self._serving_profiles.AsyncServingProfiles(self._api_client) # type: ignore[no-any-return] + async def aclose(self) -> None: """Closes the async client explicitly. @@ -399,6 +415,7 @@ def __init__( self._sessions: Optional[ModuleType] = None self._sandboxes: Optional[ModuleType] = None self._memory_banks: Optional[ModuleType] = None + self._serving_profiles: Optional[ModuleType] = None @property def evals(self) -> "evals_module.Evals": @@ -598,3 +615,12 @@ def memory_banks(self) -> "memory_banks_module.MemoryBanks": if self._memory_banks is None: self._memory_banks = importlib.import_module(".memory_banks", __package__) return self._memory_banks.MemoryBanks(self._api_client) # type: ignore[no-any-return] + + @property + def serving_profiles(self) -> "serving_profiles_module.ServingProfiles": + if self._serving_profiles is None: + self._serving_profiles = importlib.import_module( + ".serving_profiles", + __package__, + ) + return self._serving_profiles.ServingProfiles(self._api_client) # type: ignore[no-any-return] diff --git a/agentplatform/_genai/serving_profiles.py b/agentplatform/_genai/serving_profiles.py new file mode 100644 index 0000000000..7417ac1346 --- /dev/null +++ b/agentplatform/_genai/serving_profiles.py @@ -0,0 +1,1251 @@ +# Copyright 2025 Google LLC +# +# 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. +# + +# Code generated by the Google Gen AI SDK generator DO NOT EDIT. + +import json +import logging +from typing import Any, Optional, Union +from urllib.parse import urlencode + +from google.genai import _api_module +from google.genai import _common +from google.genai._common import get_value_by_path as getv +from google.genai._common import set_value_by_path as setv +from google.genai.pagers import AsyncPager, Pager + +from . import _operations_utils +from . import types + +logger = logging.getLogger("agentplatform_genai.servingprofiles") + + +def _CreateServingProfileConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + return to_object + + +def _CreateServingProfileRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["serving_profile_id"]) is not None: + setv( + to_object, + ["_query", "servingProfileId"], + getv(from_object, ["serving_profile_id"]), + ) + + if getv(from_object, ["display_name"]) is not None: + setv(to_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["scope"]) is not None: + setv(to_object, ["scope"], getv(from_object, ["scope"])) + + if getv(from_object, ["cmek_config"]) is not None: + setv(to_object, ["cmekConfig"], getv(from_object, ["cmek_config"])) + + if getv(from_object, ["config"]) is not None: + _CreateServingProfileConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +def _DeleteServingProfileRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + +def _GetServingProfileOperationParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["operation_name"]) is not None: + setv( + to_object, ["_url", "operationName"], getv(from_object, ["operation_name"]) + ) + + return to_object + + +def _GetServingProfileRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + +def _ListServingProfilesConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["page_size"]) is not None: + setv(parent_object, ["_query", "pageSize"], getv(from_object, ["page_size"])) + + if getv(from_object, ["page_token"]) is not None: + setv(parent_object, ["_query", "pageToken"], getv(from_object, ["page_token"])) + + return to_object + + +def _ListServingProfilesRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["config"]) is not None: + setv( + to_object, + ["config"], + _ListServingProfilesConfig_to_vertex( + getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _UpdateServingProfileConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["description"]) is not None: + setv(parent_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["update_mask"]) is not None: + setv( + parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"]) + ) + + return to_object + + +def _UpdateServingProfileRequestParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _UpdateServingProfileConfig_to_vertex(getv(from_object, ["config"]), to_object) + + return to_object + + +class ServingProfiles(_api_module.BaseModule): + """Class for managing ServingProfiles.""" + + def get( + self, *, name: str, config: Optional[types.GetServingProfileConfigOrDict] = None + ) -> types.ServingProfile: + """ + Gets a ServingProfile. + """ + + parameter_model = types._GetServingProfileRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfile._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _create( + self, + *, + serving_profile_id: str, + display_name: str, + scope: types.ServingProfileScope, + cmek_config: types.ServingProfileCmekConfigOrDict, + config: Optional[types.CreateServingProfileConfigOrDict] = None, + ) -> types.ServingProfileOperation: + """ + Creates a new ServingProfile. + """ + + parameter_model = types._CreateServingProfileRequestParameters( + serving_profile_id=serving_profile_id, + display_name=display_name, + scope=scope, + cmek_config=cmek_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "servingProfiles".format_map(request_url_dict) + else: + path = "servingProfiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfileOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _update( + self, + *, + name: str, + config: Optional[types.UpdateServingProfileConfigOrDict] = None, + ) -> types.ServingProfile: + """ + Updates a ServingProfile. + """ + + parameter_model = types._UpdateServingProfileRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("patch", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfile._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def _list( + self, *, config: Optional[types.ListServingProfilesConfigOrDict] = None + ) -> types.ListServingProfilesResponse: + parameter_model = types._ListServingProfilesRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListServingProfilesRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "servingProfiles".format_map(request_url_dict) + else: + path = "servingProfiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ListServingProfilesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def delete( + self, + *, + name: str, + config: Optional[types.DeleteServingProfileConfigOrDict] = None, + ) -> None: + """ + Deletes a ServingProfile. + """ + + parameter_model = types._DeleteServingProfileRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + self._api_client.request("delete", path, request_dict, http_options) + + def _get_serving_profile_operation( + self, + *, + operation_name: str, + config: Optional[types.GetServingProfileOperationConfigOrDict] = None, + ) -> types.ServingProfileOperation: + parameter_model = types._GetServingProfileOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetServingProfileOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfileOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def create( + self, + *, + display_name: str, + scope: Union[str, types.ServingProfileScope], + serving_profile_id: str, + cmek_config: Union[dict[str, Any], types.ServingProfileCmekConfig], + config: Optional[types.CreateServingProfileConfigOrDict] = None, + ) -> Union[types.ServingProfile, types.ServingProfileOperation]: + """Creates a new ServingProfile. + + Args: + display_name (str): + Required. The display name of the ServingProfile. + scope (Union[str, ServingProfileScope]): + Required. The specific API this ServingProfile applies to. + serving_profile_id (str): + Required. The ID to use for the ServingProfile, which will become + the final component of the ServingProfile's resource name. + cmek_config (Union[dict[str, Any], ServingProfileCmekConfig]): + Required. Customer-Managed Encryption Key (CMEK) configuration for + the ServingProfile. + config (CreateServingProfileConfigOrDict): + Optional. The configuration for creating the ServingProfile. + + Returns: + ServingProfile: The created ServingProfile if wait_for_completion is + True. + ServingProfileOperation: The operation for creating the ServingProfile + if wait_for_completion is False. + """ + if config is None: + config = types.CreateServingProfileConfig() + elif isinstance(config, dict): + config = types.CreateServingProfileConfig.model_validate(config) + elif not isinstance(config, types.CreateServingProfileConfig): + raise TypeError( + f"config must be a dict or CreateServingProfileConfig, but got {type(config)}." + ) + + operation = self._create( + display_name=display_name, + scope=scope, + serving_profile_id=serving_profile_id, + cmek_config=cmek_config, + config=config, + ) + + if config.wait_for_completion: + operation = _operations_utils.await_operation( + operation_name=operation.name, + get_operation_fn=self._get_serving_profile_operation, + ) + if operation.error: + raise RuntimeError( + f"Failed to create ServingProfile: {operation.error}" + ) + return self.get(name=operation.response.name) + + return operation + + def update( + self, + *, + name: str, + config: Optional[types.UpdateServingProfileConfigOrDict] = None, + ) -> types.ServingProfile: + """Updates an existing ServingProfile. + + Args: + name (str): + Required. The resource name of the ServingProfile to update. + Format: + projects/{project}/locations/{location}/servingProfiles/{serving_profile} + config (UpdateServingProfileConfigOrDict): + Optional. The configuration for updating the ServingProfile. + + Returns: + ServingProfile: The updated ServingProfile. + """ + if config is None: + config = types.UpdateServingProfileConfig() + elif isinstance(config, dict): + config = types.UpdateServingProfileConfig.model_validate(config) + elif not isinstance(config, types.UpdateServingProfileConfig): + raise TypeError( + f"config must be a dict or UpdateServingProfileConfig, but got {type(config)}." + ) + + config = config.model_copy() + + update_mask_paths = [] + if config.display_name is not None: + update_mask_paths.append("displayName") + if config.description is not None: + update_mask_paths.append("description") + + if not config.update_mask and update_mask_paths: + config.update_mask = ",".join(update_mask_paths) + + if not config.update_mask: + raise ValueError( + "At least one field to update must be provided or specified via `update_mask`." + ) + + return self._update(name=name, config=config) + + def list( + self, *, config: Optional[types.ListServingProfilesConfigOrDict] = None + ) -> Pager[types.ServingProfile]: + """ + Lists ServingProfiles in a Location. + """ + + list_request = self._list + return Pager( + "serving_profiles", + list_request, + self._list(config=config), + config, + ) + + +class AsyncServingProfiles(_api_module.BaseModule): + """Class for managing ServingProfiles.""" + + async def get( + self, *, name: str, config: Optional[types.GetServingProfileConfigOrDict] = None + ) -> types.ServingProfile: + """ + Gets a ServingProfile. + """ + + parameter_model = types._GetServingProfileRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfile._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _create( + self, + *, + serving_profile_id: str, + display_name: str, + scope: types.ServingProfileScope, + cmek_config: types.ServingProfileCmekConfigOrDict, + config: Optional[types.CreateServingProfileConfigOrDict] = None, + ) -> types.ServingProfileOperation: + """ + Creates a new ServingProfile. + """ + + parameter_model = types._CreateServingProfileRequestParameters( + serving_profile_id=serving_profile_id, + display_name=display_name, + scope=scope, + cmek_config=cmek_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "servingProfiles".format_map(request_url_dict) + else: + path = "servingProfiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfileOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _update( + self, + *, + name: str, + config: Optional[types.UpdateServingProfileConfigOrDict] = None, + ) -> types.ServingProfile: + """ + Updates a ServingProfile. + """ + + parameter_model = types._UpdateServingProfileRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _UpdateServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfile._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _list( + self, *, config: Optional[types.ListServingProfilesConfigOrDict] = None + ) -> types.ListServingProfilesResponse: + parameter_model = types._ListServingProfilesRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _ListServingProfilesRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "servingProfiles".format_map(request_url_dict) + else: + path = "servingProfiles" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ListServingProfilesResponse._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def delete( + self, + *, + name: str, + config: Optional[types.DeleteServingProfileConfigOrDict] = None, + ) -> None: + """ + Deletes a ServingProfile. + """ + + parameter_model = types._DeleteServingProfileRequestParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _DeleteServingProfileRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + await self._api_client.async_request("delete", path, request_dict, http_options) + + async def _get_serving_profile_operation( + self, + *, + operation_name: str, + config: Optional[types.GetServingProfileOperationConfigOrDict] = None, + ) -> types.ServingProfileOperation: + parameter_model = types._GetServingProfileOperationParameters( + operation_name=operation_name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetServingProfileOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.ServingProfileOperation._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def create( + self, + *, + display_name: str, + scope: Union[str, types.ServingProfileScope], + serving_profile_id: str, + cmek_config: Union[dict[str, Any], types.ServingProfileCmekConfig], + config: Optional[types.CreateServingProfileConfigOrDict] = None, + ) -> Union[types.ServingProfile, types.ServingProfileOperation]: + """Creates a new ServingProfile asynchronously. + + Args: + display_name (str): + Required. The display name of the ServingProfile. + scope (Union[str, ServingProfileScope]): + Required. The specific API this ServingProfile applies to. + serving_profile_id (str): + Required. The ID to use for the ServingProfile, which will become + the final component of the ServingProfile's resource name. + cmek_config (Union[dict[str, Any], ServingProfileCmekConfig]): + Required. Customer-Managed Encryption Key (CMEK) configuration for + the ServingProfile. + config (CreateServingProfileConfigOrDict): + Optional. The configuration for creating the ServingProfile. + + Returns: + ServingProfile: The created ServingProfile if wait_for_completion is + True. + ServingProfileOperation: The operation for creating the ServingProfile + if wait_for_completion is False. + """ + if config is None: + config = types.CreateServingProfileConfig() + elif isinstance(config, dict): + config = types.CreateServingProfileConfig.model_validate(config) + elif not isinstance(config, types.CreateServingProfileConfig): + raise TypeError( + f"config must be a dict or CreateServingProfileConfig, but got {type(config)}." + ) + + operation = await self._create( + display_name=display_name, + scope=scope, + serving_profile_id=serving_profile_id, + cmek_config=cmek_config, + config=config, + ) + + if config.wait_for_completion: + operation = await _operations_utils.await_operation_async( + operation_name=operation.name, + get_operation_fn=self._get_serving_profile_operation, + ) + if operation.error: + raise RuntimeError( + f"Failed to create ServingProfile: {operation.error}" + ) + return await self.get(name=operation.response.name) + + return operation + + async def update( + self, + *, + name: str, + config: Optional[types.UpdateServingProfileConfigOrDict] = None, + ) -> types.ServingProfile: + """Updates an existing ServingProfile asynchronously. + + Args: + name (str): + Required. The resource name of the ServingProfile to update. + Format: + projects/{project}/locations/{location}/servingProfiles/{serving_profile} + config (UpdateServingProfileConfigOrDict): + Optional. The configuration for updating the ServingProfile. + + Returns: + ServingProfile: The updated ServingProfile. + """ + if config is None: + config = types.UpdateServingProfileConfig() + elif isinstance(config, dict): + config = types.UpdateServingProfileConfig.model_validate(config) + elif not isinstance(config, types.UpdateServingProfileConfig): + raise TypeError( + f"config must be a dict or UpdateServingProfileConfig, but got {type(config)}." + ) + + config = config.model_copy() + + update_mask_paths = [] + if config.display_name is not None: + update_mask_paths.append("displayName") + if config.description is not None: + update_mask_paths.append("description") + + if not config.update_mask and update_mask_paths: + config.update_mask = ",".join(update_mask_paths) + + if not config.update_mask: + raise ValueError( + "At least one field to update must be provided or specified via `update_mask`." + ) + + return await self._update(name=name, config=config) + + async def list( + self, *, config: Optional[types.ListServingProfilesConfigOrDict] = None + ) -> AsyncPager[types.ServingProfile]: + list_request = self._list + return AsyncPager( + "serving_profiles", + list_request, + await self._list(config=config), + config, + ) diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index b2ff7bcd22..40378c1d0f 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -46,6 +46,7 @@ from .common import _CreateRuntimeSessionRequestParameters from .common import _CreateSandboxEnvironmentSnapshotRequestParameters from .common import _CreateSandboxEnvironmentTemplateRequestParameters +from .common import _CreateServingProfileRequestParameters from .common import _CreateSkillRequestParameters from .common import _CustomJobParameters from .common import _CustomJobParameters @@ -68,6 +69,7 @@ from .common import _DeleteRuntimeSessionRequestParameters from .common import _DeleteSandboxEnvironmentSnapshotRequestParameters from .common import _DeleteSandboxEnvironmentTemplateRequestParameters +from .common import _DeleteServingProfileRequestParameters from .common import _DeleteSkillRequestParameters from .common import _DeployRequestParameters from .common import _EvaluateInstancesRequestParameters @@ -126,6 +128,8 @@ from .common import _GetSandboxEnvironmentSnapshotRequestParameters from .common import _GetSandboxEnvironmentTemplateOperationParameters from .common import _GetSandboxEnvironmentTemplateRequestParameters +from .common import _GetServingProfileOperationParameters +from .common import _GetServingProfileRequestParameters from .common import _GetSkillOperationParameters from .common import _GetSkillRequestParameters from .common import _GetSkillRevisionRequestParameters @@ -152,6 +156,7 @@ from .common import _ListRuntimeSessionsRequestParameters from .common import _ListSandboxEnvironmentSnapshotsRequestParameters from .common import _ListSandboxEnvironmentTemplatesRequestParameters +from .common import _ListServingProfilesRequestParameters from .common import _ListSkillRevisionsRequestParameters from .common import _ListSkillsRequestParameters from .common import _OptimizeRequestParameters @@ -186,6 +191,7 @@ from .common import _UpdateRuntimeFeedbackEntryRequestParameters from .common import _UpdateRuntimeRequestParameters from .common import _UpdateRuntimeSessionRequestParameters +from .common import _UpdateServingProfileRequestParameters from .common import _UpdateSkillRequestParameters from .common import _UploadRagFileParameters from .common import _UpsertExamplesParameters @@ -382,6 +388,9 @@ from .common import CreateSandboxEnvironmentTemplateConfig from .common import CreateSandboxEnvironmentTemplateConfigDict from .common import CreateSandboxEnvironmentTemplateConfigOrDict +from .common import CreateServingProfileConfig +from .common import CreateServingProfileConfigDict +from .common import CreateServingProfileConfigOrDict from .common import CreateSkillConfig from .common import CreateSkillConfigDict from .common import CreateSkillConfigOrDict @@ -516,6 +525,9 @@ from .common import DeleteSandboxEnvironmentTemplateOperation from .common import DeleteSandboxEnvironmentTemplateOperationDict from .common import DeleteSandboxEnvironmentTemplateOperationOrDict +from .common import DeleteServingProfileConfig +from .common import DeleteServingProfileConfigDict +from .common import DeleteServingProfileConfigOrDict from .common import DeleteSkillConfig from .common import DeleteSkillConfigDict from .common import DeleteSkillConfigOrDict @@ -965,6 +977,12 @@ from .common import GetSandboxEnvironmentTemplateConfig from .common import GetSandboxEnvironmentTemplateConfigDict from .common import GetSandboxEnvironmentTemplateConfigOrDict +from .common import GetServingProfileConfig +from .common import GetServingProfileConfigDict +from .common import GetServingProfileConfigOrDict +from .common import GetServingProfileOperationConfig +from .common import GetServingProfileOperationConfigDict +from .common import GetServingProfileOperationConfigOrDict from .common import GetSkillConfig from .common import GetSkillConfigDict from .common import GetSkillConfigOrDict @@ -1175,6 +1193,12 @@ from .common import ListSandboxEnvironmentTemplatesResponse from .common import ListSandboxEnvironmentTemplatesResponseDict from .common import ListSandboxEnvironmentTemplatesResponseOrDict +from .common import ListServingProfilesConfig +from .common import ListServingProfilesConfigDict +from .common import ListServingProfilesConfigOrDict +from .common import ListServingProfilesResponse +from .common import ListServingProfilesResponseDict +from .common import ListServingProfilesResponseOrDict from .common import ListSkillRevisionsConfig from .common import ListSkillRevisionsConfigDict from .common import ListSkillRevisionsConfigOrDict @@ -2069,6 +2093,16 @@ from .common import SecretRef from .common import SecretRefDict from .common import SecretRefOrDict +from .common import ServingProfile +from .common import ServingProfileCmekConfig +from .common import ServingProfileCmekConfigDict +from .common import ServingProfileCmekConfigOrDict +from .common import ServingProfileDict +from .common import ServingProfileOperation +from .common import ServingProfileOperationDict +from .common import ServingProfileOperationOrDict +from .common import ServingProfileOrDict +from .common import ServingProfileScope from .common import Session from .common import SessionDict from .common import SessionEvent @@ -2258,6 +2292,9 @@ from .common import UpdateRuntimeSessionConfig from .common import UpdateRuntimeSessionConfigDict from .common import UpdateRuntimeSessionConfigOrDict +from .common import UpdateServingProfileConfig +from .common import UpdateServingProfileConfigDict +from .common import UpdateServingProfileConfigOrDict from .common import UpdateSkillConfig from .common import UpdateSkillConfigDict from .common import UpdateSkillConfigOrDict @@ -4196,6 +4233,36 @@ "GetExampleStoreOperationConfig", "GetExampleStoreOperationConfigDict", "GetExampleStoreOperationConfigOrDict", + "GetServingProfileConfig", + "GetServingProfileConfigDict", + "GetServingProfileConfigOrDict", + "ServingProfileCmekConfig", + "ServingProfileCmekConfigDict", + "ServingProfileCmekConfigOrDict", + "ServingProfile", + "ServingProfileDict", + "ServingProfileOrDict", + "CreateServingProfileConfig", + "CreateServingProfileConfigDict", + "CreateServingProfileConfigOrDict", + "ServingProfileOperation", + "ServingProfileOperationDict", + "ServingProfileOperationOrDict", + "UpdateServingProfileConfig", + "UpdateServingProfileConfigDict", + "UpdateServingProfileConfigOrDict", + "ListServingProfilesConfig", + "ListServingProfilesConfigDict", + "ListServingProfilesConfigOrDict", + "ListServingProfilesResponse", + "ListServingProfilesResponseDict", + "ListServingProfilesResponseOrDict", + "DeleteServingProfileConfig", + "DeleteServingProfileConfigDict", + "DeleteServingProfileConfigOrDict", + "GetServingProfileOperationConfig", + "GetServingProfileOperationConfigDict", + "GetServingProfileOperationConfigOrDict", "PromptOptimizerConfig", "PromptOptimizerConfigDict", "PromptOptimizerConfigOrDict", @@ -4344,6 +4411,7 @@ "FunctionResponseScheduling", "MediaResolution", "ArrayOperator", + "ServingProfileScope", "EvaluationExperimentMergeStrategy", "EvaluationItemType", "SamplingMethod", @@ -4542,6 +4610,12 @@ "_FetchExamplesParameters", "_RemoveExamplesParameters", "_GetExampleStoreOperationParameters", + "_GetServingProfileRequestParameters", + "_CreateServingProfileRequestParameters", + "_UpdateServingProfileRequestParameters", + "_ListServingProfilesRequestParameters", + "_DeleteServingProfileRequestParameters", + "_GetServingProfileOperationParameters", "evals", "agent_engines", "prompts", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 7c83714a17..4c3da8b398 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -659,6 +659,19 @@ class ArrayOperator(_common.CaseInSensitiveEnum): """The metadata array field in the example must contain all of the values.""" +class ServingProfileScope(_common.CaseInSensitiveEnum): + """The specific API this ServingProfile applies to.""" + + SERVING_PROFILE_SCOPE_UNSPECIFIED = "SERVING_PROFILE_SCOPE_UNSPECIFIED" + """Unspecified scope.""" + GEMINI_LIVE = "GEMINI_LIVE" + """Scope for Gemini Live.""" + INTERACTIONS_API = "INTERACTIONS_API" + """Scope for Interactions API.""" + RESPONSE_API = "RESPONSE_API" + """Scope for Response API.""" + + class EvaluationExperimentMergeStrategy(_common.CaseInSensitiveEnum): """Merge strategy for the evaluation experiment.""" @@ -29497,6 +29510,488 @@ class _GetExampleStoreOperationParametersDict(TypedDict, total=False): ] +class GetServingProfileConfig(_common.BaseModel): + """Config for getting a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetServingProfileConfigDict(TypedDict, total=False): + """Config for getting a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetServingProfileConfigOrDict = Union[ + GetServingProfileConfig, GetServingProfileConfigDict +] + + +class _GetServingProfileRequestParameters(_common.BaseModel): + """Parameters for GetServingProfileRequest.""" + + name: Optional[str] = Field( + default=None, description="""The resource name of the ServingProfile.""" + ) + config: Optional[GetServingProfileConfig] = Field(default=None, description="""""") + + +class _GetServingProfileRequestParametersDict(TypedDict, total=False): + """Parameters for GetServingProfileRequest.""" + + name: Optional[str] + """The resource name of the ServingProfile.""" + + config: Optional[GetServingProfileConfigDict] + """""" + + +_GetServingProfileRequestParametersOrDict = Union[ + _GetServingProfileRequestParameters, _GetServingProfileRequestParametersDict +] + + +class ServingProfileCmekConfig(_common.BaseModel): + """Configuration for Customer-Managed Encryption Keys (CMEK).""" + + encryption_spec: Optional[genai_types.EncryptionSpec] = Field( + default=None, + description="""Required. The customer-managed encryption key spec for the Serving Profile.""", + ) + + +class ServingProfileCmekConfigDict(TypedDict, total=False): + """Configuration for Customer-Managed Encryption Keys (CMEK).""" + + encryption_spec: Optional[genai_types.EncryptionSpecDict] + """Required. The customer-managed encryption key spec for the Serving Profile.""" + + +ServingProfileCmekConfigOrDict = Union[ + ServingProfileCmekConfig, ServingProfileCmekConfigDict +] + + +class ServingProfile(_common.BaseModel): + """Represents a ServingProfile resource.""" + + scope: Optional[ServingProfileScope] = Field( + default=None, + description="""Required. The specific API this ServingProfile applies to.""", + ) + cmek_config: Optional[ServingProfileCmekConfig] = Field( + default=None, description="""CMEK configuration for the ServingProfile.""" + ) + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when the ServingProfile was created.""", + ) + description: Optional[str] = Field( + default=None, description="""Optional. The description of the ServingProfile.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the ServingProfile. The name can be up to 128 characters long and can consist of any UTF-8 characters.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the ServingProfile.""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when the ServingProfile was last updated.""", + ) + + +class ServingProfileDict(TypedDict, total=False): + """Represents a ServingProfile resource.""" + + scope: Optional[ServingProfileScope] + """Required. The specific API this ServingProfile applies to.""" + + cmek_config: Optional[ServingProfileCmekConfigDict] + """CMEK configuration for the ServingProfile.""" + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when the ServingProfile was created.""" + + description: Optional[str] + """Optional. The description of the ServingProfile.""" + + display_name: Optional[str] + """Required. The display name of the ServingProfile. The name can be up to 128 characters long and can consist of any UTF-8 characters.""" + + name: Optional[str] + """Identifier. The resource name of the ServingProfile.""" + + update_time: Optional[datetime.datetime] + """Output only. Timestamp when the ServingProfile was last updated.""" + + +ServingProfileOrDict = Union[ServingProfile, ServingProfileDict] + + +class CreateServingProfileConfig(_common.BaseModel): + """Config for creating a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to wait for the creation LRO to complete.""", + ) + description: Optional[str] = Field( + default=None, description="""Optional. The description.""" + ) + + +class CreateServingProfileConfigDict(TypedDict, total=False): + """Config for creating a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + wait_for_completion: Optional[bool] + """Whether to wait for the creation LRO to complete.""" + + description: Optional[str] + """Optional. The description.""" + + +CreateServingProfileConfigOrDict = Union[ + CreateServingProfileConfig, CreateServingProfileConfigDict +] + + +class _CreateServingProfileRequestParameters(_common.BaseModel): + """Parameters for creating a serving profile.""" + + serving_profile_id: Optional[str] = Field( + default=None, description="""Required. The ID to use for the ServingProfile.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the ServingProfile.""", + ) + scope: Optional[ServingProfileScope] = Field( + default=None, + description="""Required. The specific API this ServingProfile applies to.""", + ) + cmek_config: Optional[ServingProfileCmekConfig] = Field( + default=None, description="""Required. CMEK configuration.""" + ) + config: Optional[CreateServingProfileConfig] = Field( + default=None, description="""""" + ) + + +class _CreateServingProfileRequestParametersDict(TypedDict, total=False): + """Parameters for creating a serving profile.""" + + serving_profile_id: Optional[str] + """Required. The ID to use for the ServingProfile.""" + + display_name: Optional[str] + """Required. The display name of the ServingProfile.""" + + scope: Optional[ServingProfileScope] + """Required. The specific API this ServingProfile applies to.""" + + cmek_config: Optional[ServingProfileCmekConfigDict] + """Required. CMEK configuration.""" + + config: Optional[CreateServingProfileConfigDict] + """""" + + +_CreateServingProfileRequestParametersOrDict = Union[ + _CreateServingProfileRequestParameters, _CreateServingProfileRequestParametersDict +] + + +class ServingProfileOperation(_common.BaseModel): + """Operation that has a serving profile as a response.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[ServingProfile] = Field( + default=None, description="""The created ServingProfile.""" + ) + + +class ServingProfileOperationDict(TypedDict, total=False): + """Operation that has a serving profile as a response.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + response: Optional[ServingProfileDict] + """The created ServingProfile.""" + + +ServingProfileOperationOrDict = Union[ + ServingProfileOperation, ServingProfileOperationDict +] + + +class UpdateServingProfileConfig(_common.BaseModel): + """Config for updating a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field(default=None, description="""""") + description: Optional[str] = Field(default=None, description="""""") + update_mask: Optional[str] = Field( + default=None, + description="""Update mask. If not provided, it will be inferred based on the other fields provided to the update method. If provided, only the fields that are specified in the update mask will be updated.""", + ) + + +class UpdateServingProfileConfigDict(TypedDict, total=False): + """Config for updating a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + display_name: Optional[str] + """""" + + description: Optional[str] + """""" + + update_mask: Optional[str] + """Update mask. If not provided, it will be inferred based on the other fields provided to the update method. If provided, only the fields that are specified in the update mask will be updated.""" + + +UpdateServingProfileConfigOrDict = Union[ + UpdateServingProfileConfig, UpdateServingProfileConfigDict +] + + +class _UpdateServingProfileRequestParameters(_common.BaseModel): + """Parameters for updating a serving profile.""" + + name: Optional[str] = Field( + default=None, description="""Required. The resource name.""" + ) + config: Optional[UpdateServingProfileConfig] = Field( + default=None, description="""""" + ) + + +class _UpdateServingProfileRequestParametersDict(TypedDict, total=False): + """Parameters for updating a serving profile.""" + + name: Optional[str] + """Required. The resource name.""" + + config: Optional[UpdateServingProfileConfigDict] + """""" + + +_UpdateServingProfileRequestParametersOrDict = Union[ + _UpdateServingProfileRequestParameters, _UpdateServingProfileRequestParametersDict +] + + +class ListServingProfilesConfig(_common.BaseModel): + """Config for listing ServingProfiles.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + page_size: Optional[int] = Field(default=None, description="""""") + page_token: Optional[str] = Field(default=None, description="""""") + + +class ListServingProfilesConfigDict(TypedDict, total=False): + """Config for listing ServingProfiles.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + +ListServingProfilesConfigOrDict = Union[ + ListServingProfilesConfig, ListServingProfilesConfigDict +] + + +class _ListServingProfilesRequestParameters(_common.BaseModel): + """Parameters for listing ServingProfiles.""" + + config: Optional[ListServingProfilesConfig] = Field( + default=None, description="""""" + ) + + +class _ListServingProfilesRequestParametersDict(TypedDict, total=False): + """Parameters for listing ServingProfiles.""" + + config: Optional[ListServingProfilesConfigDict] + """""" + + +_ListServingProfilesRequestParametersOrDict = Union[ + _ListServingProfilesRequestParameters, _ListServingProfilesRequestParametersDict +] + + +class ListServingProfilesResponse(_common.BaseModel): + """Response for listing ServingProfiles.""" + + sdk_http_response: Optional[genai_types.HttpResponse] = Field( + default=None, description="""Used to retain the full HTTP response.""" + ) + next_page_token: Optional[str] = Field(default=None, description="""""") + serving_profiles: Optional[list[ServingProfile]] = Field( + default=None, description="""List of ServingProfiles.""" + ) + + +class ListServingProfilesResponseDict(TypedDict, total=False): + """Response for listing ServingProfiles.""" + + sdk_http_response: Optional[genai_types.HttpResponse] + """Used to retain the full HTTP response.""" + + next_page_token: Optional[str] + """""" + + serving_profiles: Optional[list[ServingProfileDict]] + """List of ServingProfiles.""" + + +ListServingProfilesResponseOrDict = Union[ + ListServingProfilesResponse, ListServingProfilesResponseDict +] + + +class DeleteServingProfileConfig(_common.BaseModel): + """Config for deleting a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class DeleteServingProfileConfigDict(TypedDict, total=False): + """Config for deleting a serving profile.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +DeleteServingProfileConfigOrDict = Union[ + DeleteServingProfileConfig, DeleteServingProfileConfigDict +] + + +class _DeleteServingProfileRequestParameters(_common.BaseModel): + """Parameters for deleting a serving profile.""" + + name: Optional[str] = Field( + default=None, description="""Required. The resource name.""" + ) + config: Optional[DeleteServingProfileConfig] = Field( + default=None, description="""""" + ) + + +class _DeleteServingProfileRequestParametersDict(TypedDict, total=False): + """Parameters for deleting a serving profile.""" + + name: Optional[str] + """Required. The resource name.""" + + config: Optional[DeleteServingProfileConfigDict] + """""" + + +_DeleteServingProfileRequestParametersOrDict = Union[ + _DeleteServingProfileRequestParameters, _DeleteServingProfileRequestParametersDict +] + + +class GetServingProfileOperationConfig(_common.BaseModel): + """Config for getting a serving profile operation.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetServingProfileOperationConfigDict(TypedDict, total=False): + """Config for getting a serving profile operation.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetServingProfileOperationConfigOrDict = Union[ + GetServingProfileOperationConfig, GetServingProfileOperationConfigDict +] + + +class _GetServingProfileOperationParameters(_common.BaseModel): + """Parameters for getting a serving profile operation.""" + + operation_name: Optional[str] = Field( + default=None, description="""The operation name.""" + ) + config: Optional[GetServingProfileOperationConfig] = Field( + default=None, description="""""" + ) + + +class _GetServingProfileOperationParametersDict(TypedDict, total=False): + """Parameters for getting a serving profile operation.""" + + operation_name: Optional[str] + """The operation name.""" + + config: Optional[GetServingProfileOperationConfigDict] + """""" + + +_GetServingProfileOperationParametersOrDict = Union[ + _GetServingProfileOperationParameters, _GetServingProfileOperationParametersDict +] + + class PromptOptimizerConfig(_common.BaseModel): """VAPO Prompt Optimizer Config.""" diff --git a/tests/unit/agentplatform/genai/replays/test_serving_profiles_create.py b/tests/unit/agentplatform/genai/replays/test_serving_profiles_create.py new file mode 100644 index 0000000000..8078571461 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_serving_profiles_create.py @@ -0,0 +1,54 @@ +"""Tests the serving_profiles.create() method against the Vertex AI endpoint using replays.""" + +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper +from google.genai import types as genai_types + +# MANDATORY: Initialize the replay test framework for this module +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_create_serving_profile(client): + """Tests synchronous serving profile creation with LRO polling.""" + location = "us" + client._api_client.location = location + client._api_client._http_options.base_url = ( + f"https://aiplatform.{location}.rep.googleapis.com/" + ) + kms_key_name = ( + f"projects/ucaip-e2e-test-kms-key-host/locations/{location}/" + f"keyRings/{location}-keys/cryptoKeys/e2e-test-key" + ) + cmek_config = types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec(kms_key_name=kms_key_name) + ) + + sp = None + test_display_name = "SDK Replay Test Profile" + test_description = "Created by SDK Replay Test" + test_scope = types.ServingProfileScope.GEMINI_LIVE + try: + sp = client.serving_profiles.create( + serving_profile_id="sp-sdk-replay-test", + display_name=test_display_name, + scope=test_scope, + cmek_config=cmek_config, + config=types.CreateServingProfileConfig(description=test_description), + ) + + assert sp.name is not None + assert sp.display_name == test_display_name + assert sp.description == test_description + assert sp.scope == test_scope + assert sp.cmek_config is not None + assert sp.cmek_config.encryption_spec.kms_key_name == kms_key_name + assert sp.create_time is not None + finally: + if sp: + try: + client.serving_profiles.delete(name=sp.name) + except Exception: + pass diff --git a/tests/unit/agentplatform/genai/replays/test_serving_profiles_delete.py b/tests/unit/agentplatform/genai/replays/test_serving_profiles_delete.py new file mode 100644 index 0000000000..2c4417c6cd --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_serving_profiles_delete.py @@ -0,0 +1,51 @@ +"""Tests the serving_profiles.delete() method against the Vertex AI endpoint using replays.""" + +from google.api_core import exceptions +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper +from google.genai import errors +from google.genai import types as genai_types +import pytest + +# MANDATORY: Initialize the replay test framework for this module +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_delete_serving_profile(client): + """Tests deleting a serving profile.""" + location = "us" + client._api_client.location = location + client._api_client._http_options.base_url = ( + f"https://aiplatform.{location}.rep.googleapis.com/" + ) + kms_key_name = ( + f"projects/ucaip-e2e-test-kms-key-host/locations/{location}/" + f"keyRings/{location}-keys/cryptoKeys/e2e-test-key" + ) + cmek_config = types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec(kms_key_name=kms_key_name) + ) + + try: + sp = client.serving_profiles.create( + serving_profile_id="sp-sdk-replay-test", + display_name="SDK Replay Test Profile", + scope=types.ServingProfileScope.GEMINI_LIVE, + cmek_config=cmek_config, + config=types.CreateServingProfileConfig( + description="Created by SDK Replay Test" + ), + ) + + # Delete the test profile + client.serving_profiles.delete(name=sp.name) + + # Confirm deletion by fetching (should raise a 404 ClientError) + with pytest.raises(errors.ClientError, match="404"): + client.serving_profiles.get(name=sp.name) + + except exceptions.GoogleAPIError as e: + pytest.fail(f"Error calling client.serving_profiles.delete(): {e}") diff --git a/tests/unit/agentplatform/genai/replays/test_serving_profiles_get.py b/tests/unit/agentplatform/genai/replays/test_serving_profiles_get.py new file mode 100644 index 0000000000..1baca0564b --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_serving_profiles_get.py @@ -0,0 +1,62 @@ +"""Tests the serving_profiles.get() method against the Vertex AI endpoint using replays.""" + +from google.api_core import exceptions +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper +from google.genai import types as genai_types +import pytest + +# MANDATORY: Initialize the replay test framework for this module +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_get_serving_profile(client): + """Tests fetching a serving profile by name.""" + location = "us" + client._api_client.location = location + client._api_client._http_options.base_url = ( + f"https://aiplatform.{location}.rep.googleapis.com/" + ) + kms_key_name = ( + f"projects/ucaip-e2e-test-kms-key-host/locations/{location}/" + f"keyRings/{location}-keys/cryptoKeys/e2e-test-key" + ) + cmek_config = types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec(kms_key_name=kms_key_name) + ) + + sp = None + test_display_name = "SDK Replay Test Profile" + test_description = "Created by SDK Replay Test" + test_scope = types.ServingProfileScope.GEMINI_LIVE + try: + # Create a fresh profile for the get test + sp = client.serving_profiles.create( + serving_profile_id="sp-sdk-replay-test", + display_name=test_display_name, + scope=test_scope, + cmek_config=cmek_config, + config=types.CreateServingProfileConfig(description=test_description), + ) + + # Act: Get the profile + fetched_sp = client.serving_profiles.get(name=sp.name) + assert fetched_sp.name == sp.name + assert fetched_sp.display_name == test_display_name + assert fetched_sp.description == test_description + assert fetched_sp.scope == test_scope + assert fetched_sp.cmek_config is not None + assert fetched_sp.cmek_config.encryption_spec.kms_key_name == kms_key_name + assert fetched_sp.create_time is not None + + except exceptions.GoogleAPIError as e: + pytest.fail(f"Error calling client.serving_profiles.get(): {e}") + finally: + if sp: + try: + client.serving_profiles.delete(name=sp.name) + except Exception: + pass diff --git a/tests/unit/agentplatform/genai/replays/test_serving_profiles_list.py b/tests/unit/agentplatform/genai/replays/test_serving_profiles_list.py new file mode 100644 index 0000000000..c09a3ffb78 --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_serving_profiles_list.py @@ -0,0 +1,63 @@ +"""Tests the serving_profiles.list() method against the Vertex AI endpoint using replays.""" + +from google.api_core import exceptions +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper +from google.genai import types as genai_types +import pytest + +# MANDATORY: Initialize the replay test framework for this module +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_list_serving_profiles(client): + """Tests listing serving profiles via the unified pagination mechanism.""" + location = "us" + client._api_client.location = location + client._api_client._http_options.base_url = ( + f"https://aiplatform.{location}.rep.googleapis.com/" + ) + kms_key_name = ( + f"projects/ucaip-e2e-test-kms-key-host/locations/{location}/" + f"keyRings/{location}-keys/cryptoKeys/e2e-test-key" + ) + cmek_config = types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec(kms_key_name=kms_key_name) + ) + + sp = None + try: + # Guarantee there is at least one profile to list + sp = client.serving_profiles.create( + serving_profile_id="sp-sdk-replay-test", + display_name="SDK Replay Test Profile", + scope=types.ServingProfileScope.GEMINI_LIVE, + cmek_config=cmek_config, + config=types.CreateServingProfileConfig( + description="Created by SDK Replay Test" + ), + ) + + pager = client.serving_profiles.list() + + # Verify list items parse successfully + for profile in pager: + assert profile.name is not None + assert hasattr(profile, "display_name") + assert hasattr(profile, "description") + assert hasattr(profile, "scope") + assert hasattr(profile, "cmek_config") + assert hasattr(profile, "create_time") + assert hasattr(profile, "update_time") + + except exceptions.GoogleAPIError as e: + pytest.fail(f"Error calling client.serving_profiles.list(): {e}") + finally: + if sp: + try: + client.serving_profiles.delete(name=sp.name) + except Exception: + pass diff --git a/tests/unit/agentplatform/genai/replays/test_serving_profiles_update.py b/tests/unit/agentplatform/genai/replays/test_serving_profiles_update.py new file mode 100644 index 0000000000..b3e8545bef --- /dev/null +++ b/tests/unit/agentplatform/genai/replays/test_serving_profiles_update.py @@ -0,0 +1,118 @@ +"""Tests the serving_profiles.update() method against the Vertex AI endpoint using replays.""" + +from google.api_core import exceptions +from agentplatform._genai import types +from tests.unit.agentplatform.genai.replays import pytest_helper +from google.genai import types as genai_types +import pytest + +# MANDATORY: Initialize the replay test framework for this module +pytestmark = pytest_helper.setup( + file=__file__, + globals_for_file=globals(), +) + + +def test_update_serving_profile(client): + """Tests updating a serving profile.""" + location = "us" + client._api_client.location = location + client._api_client._http_options.base_url = ( + f"https://aiplatform.{location}.rep.googleapis.com/" + ) + kms_key_name = ( + f"projects/ucaip-e2e-test-kms-key-host/locations/{location}/" + f"keyRings/{location}-keys/cryptoKeys/e2e-test-key" + ) + cmek_config = types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec(kms_key_name=kms_key_name) + ) + + sp = None + test_display_name = "SDK Replay Test Profile" + test_updated_display_name = "Updated SDK Replay Test Profile" + test_description = "Created by SDK Replay Test" + test_updated_description = "Updated by SDK Replay Test" + try: + sp = client.serving_profiles.create( + serving_profile_id="sp-sdk-replay-test", + display_name=test_display_name, + scope=types.ServingProfileScope.GEMINI_LIVE, + cmek_config=cmek_config, + config=types.CreateServingProfileConfig(description=test_description), + ) + assert sp.display_name == test_display_name + assert sp.description == test_description + + updated_sp = client.serving_profiles.update( + name=sp.name, + config=types.UpdateServingProfileConfig( + display_name=test_updated_display_name, + description=test_updated_description, + ), + ) + assert updated_sp.name == sp.name + assert updated_sp.display_name == test_updated_display_name + assert updated_sp.description == test_updated_description + + except exceptions.GoogleAPIError as e: + pytest.fail(f"Error calling client.serving_profiles.update(): {e}") + finally: + if sp: + try: + client.serving_profiles.delete(name=sp.name) + except Exception: + pass + + +def test_update_serving_profile_explicit_mask(client): + """Tests updating a serving profile using explicit mask.""" + location = "us" + client._api_client.location = location + client._api_client._http_options.base_url = ( + f"https://aiplatform.{location}.rep.googleapis.com/" + ) + kms_key_name = ( + f"projects/ucaip-e2e-test-kms-key-host/locations/{location}/" + f"keyRings/{location}-keys/cryptoKeys/e2e-test-key" + ) + cmek_config = types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec(kms_key_name=kms_key_name) + ) + + sp = None + test_display_name = "SDK Replay Test Profile" + test_description = "Created by SDK Replay Test" + test_updated_description = "Updated by SDK Replay Test" + try: + sp = client.serving_profiles.create( + serving_profile_id="sp-sdk-replay-test", + display_name=test_display_name, + scope=types.ServingProfileScope.GEMINI_LIVE, + cmek_config=cmek_config, + config=types.CreateServingProfileConfig(description=test_description), + ) + assert sp.display_name == test_display_name + assert sp.description == test_description + + update_config = types.UpdateServingProfileConfig( + display_name="Display name should not change to this", + description=test_updated_description, + update_mask="description", + ) + updated_sp = client.serving_profiles.update( + name=sp.name, + config=update_config, + ) + assert updated_sp.name == sp.name + assert updated_sp.display_name == test_display_name + assert updated_sp.description == test_updated_description + + except exceptions.GoogleAPIError as e: + pytest.fail(f"Error calling client.serving_profiles.update(): {e}") + finally: + if sp: + try: + client.serving_profiles.delete(name=sp.name) + except Exception: + pass diff --git a/tests/unit/agentplatform/genai/test_serving_profiles.py b/tests/unit/agentplatform/genai/test_serving_profiles.py new file mode 100644 index 0000000000..adee8e8bbd --- /dev/null +++ b/tests/unit/agentplatform/genai/test_serving_profiles.py @@ -0,0 +1,233 @@ +# Copyright 2026 Google LLC +# +# 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 +from unittest import mock + +import google.auth.credentials +from agentplatform._genai import client as agentplatform_client +from agentplatform._genai import types as agentplatform_types +from google.genai import types as genai_types +import pytest + + +@pytest.fixture +def serving_profiles_client(): + creds = mock.create_autospec(google.auth.credentials.Credentials, instance=True) + creds.token = "test_token" + client = agentplatform_client.Client( + project="test-project", location="us", credentials=creds + ) + return client.serving_profiles + + +@pytest.fixture +def async_serving_profiles_client(): + creds = mock.create_autospec(google.auth.credentials.Credentials, instance=True) + creds.token = "test_token" + client = agentplatform_client.Client( + project="test-project", location="us", credentials=creds + ) + return client.aio.serving_profiles + + +class TestGenaiServingProfiles: + + mock_get_serving_profile_response = { + "name": ("projects/test-project/locations/us/servingProfiles/test-profile"), + "displayName": "My Test Profile", + "scope": "INTERACTIONS_API", + } + + mock_operation_response = { + "name": "projects/test-project/locations/us/operations/123", + "done": True, + "response": mock_get_serving_profile_response, + } + + def test_get_serving_profile(self, serving_profiles_client): + with mock.patch.object( + serving_profiles_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse( + body=json.dumps(self.mock_get_serving_profile_response) + ) + profile_name = ( + "projects/test-project/locations/us/servingProfiles/test-profile" + ) + profile = serving_profiles_client.get(name=profile_name) + request_mock.assert_called_once_with( + "get", + profile_name, + {"_url": {"name": profile_name}}, + None, + ) + assert isinstance(profile, agentplatform_types.ServingProfile) + assert profile.name == profile_name + assert profile.display_name == "My Test Profile" + + def test_create_serving_profile_wait(self, serving_profiles_client): + with mock.patch.object( + serving_profiles_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + # 1. return operation from _create + genai_types.HttpResponse( + body=json.dumps( + { + "name": ( + "projects/test-project/locations/us/operations/123" + ), + "done": False, + } + ) + ), + # 2. return operation from get_operation + genai_types.HttpResponse(body=json.dumps(self.mock_operation_response)), + # 3. return the actual profile from get + genai_types.HttpResponse( + body=json.dumps(self.mock_get_serving_profile_response) + ), + ] + cmek_config = agentplatform_types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec( + kms_key_name="projects/test-project/locations/us/keyRings/my-ring/cryptoKeys/my-key" + ) + ) + profile = serving_profiles_client.create( + display_name="My Test Profile", + scope="INTERACTIONS_API", + serving_profile_id="test-profile", + cmek_config=cmek_config, + ) + assert isinstance(profile, agentplatform_types.ServingProfile) + assert ( + profile.name + == "projects/test-project/locations/us/servingProfiles/test-profile" + ) + + def test_update_serving_profile(self, serving_profiles_client): + with mock.patch.object( + serving_profiles_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse( + body=json.dumps(self.mock_get_serving_profile_response) + ) + profile_name = ( + "projects/test-project/locations/us/servingProfiles/test-profile" + ) + profile = serving_profiles_client.update( + name=profile_name, + config=agentplatform_types.UpdateServingProfileConfig( + display_name="Updated Profile", + description="New description", + ), + ) + request_mock.assert_called_once() + assert isinstance(profile, agentplatform_types.ServingProfile) + + def test_delete_serving_profile(self, serving_profiles_client): + with mock.patch.object( + serving_profiles_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse(body="{}") + profile_name = ( + "projects/test-project/locations/us/servingProfiles/test-profile" + ) + serving_profiles_client.delete(name=profile_name) + request_mock.assert_called_once() + + +class TestAsyncGenaiServingProfiles: + + mock_get_serving_profile_response = { + "name": ("projects/test-project/locations/us/servingProfiles/test-profile"), + "displayName": "My Test Profile", + "scope": "INTERACTIONS_API", + } + + mock_operation_response = { + "name": "projects/test-project/locations/us/operations/123", + "done": True, + "response": mock_get_serving_profile_response, + } + + @pytest.mark.asyncio + async def test_get_serving_profile(self, async_serving_profiles_client): + with mock.patch.object( + async_serving_profiles_client._api_client, + "async_request", + new_callable=mock.AsyncMock, + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse( + body=json.dumps(self.mock_get_serving_profile_response) + ) + profile_name = ( + "projects/test-project/locations/us/servingProfiles/test-profile" + ) + profile = await async_serving_profiles_client.get(name=profile_name) + request_mock.assert_called_once_with( + "get", + profile_name, + {"_url": {"name": profile_name}}, + None, + ) + assert isinstance(profile, agentplatform_types.ServingProfile) + assert profile.name == profile_name + assert profile.display_name == "My Test Profile" + + @pytest.mark.asyncio + async def test_create_serving_profile_wait(self, async_serving_profiles_client): + with mock.patch.object( + async_serving_profiles_client._api_client, + "async_request", + new_callable=mock.AsyncMock, + ) as request_mock: + request_mock.side_effect = [ + # 1. return operation from _create + genai_types.HttpResponse( + body=json.dumps( + { + "name": ( + "projects/test-project/locations/us/operations/123" + ), + "done": False, + } + ) + ), + # 2. return operation from get_operation + genai_types.HttpResponse(body=json.dumps(self.mock_operation_response)), + # 3. return the actual profile from get + genai_types.HttpResponse( + body=json.dumps(self.mock_get_serving_profile_response) + ), + ] + cmek_config = agentplatform_types.ServingProfileCmekConfig( + encryption_spec=genai_types.EncryptionSpec( + kms_key_name="projects/test-project/locations/us/keyRings/my-ring/cryptoKeys/my-key" + ) + ) + + with mock.patch("asyncio.sleep", new_callable=mock.AsyncMock): + profile = await async_serving_profiles_client.create( + display_name="My Test Profile", + scope="INTERACTIONS_API", + serving_profile_id="test-profile", + cmek_config=cmek_config, + ) + assert isinstance(profile, agentplatform_types.ServingProfile) + assert ( + profile.name + == "projects/test-project/locations/us/servingProfiles/test-profile" + )