diff --git a/tests/unit/vertexai/genai/test_agent_engines.py b/tests/unit/vertexai/genai/test_agent_engines.py new file mode 100644 index 0000000000..07a60fa903 --- /dev/null +++ b/tests/unit/vertexai/genai/test_agent_engines.py @@ -0,0 +1,44 @@ +# 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. +# +# pylint: disable=protected-access,bad-continuation + +import importlib +from unittest import mock + +from google.cloud import aiplatform +import vertexai +from google.cloud.aiplatform import initializer as aiplatform_initializer +from vertexai import _genai + +import pytest + +_TEST_PROJECT = "test-project" +_TEST_LOCATION = "us-central1" + + +pytestmark = pytest.mark.usefixtures("google_auth_mock") + + +class TestAgentEngines: + """Unit tests for the Agent Engines client.""" + + def setup_method(self): + importlib.reload(aiplatform_initializer) + importlib.reload(aiplatform) + importlib.reload(vertexai) + vertexai.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + ) diff --git a/vertexai/_genai/_transformers.py b/vertexai/_genai/_transformers.py index b87820d655..86bb2f33f6 100644 --- a/vertexai/_genai/_transformers.py +++ b/vertexai/_genai/_transformers.py @@ -14,3 +14,7 @@ # """Transformers module for Vertex addons.""" + +from google.genai import _transformers as genai_transformers + +t_bytes = genai_transformers.t_bytes diff --git a/vertexai/_genai/agent_engines.py b/vertexai/_genai/agent_engines.py new file mode 100644 index 0000000000..a740479f0c --- /dev/null +++ b/vertexai/_genai/agent_engines.py @@ -0,0 +1,1558 @@ +# 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 logging +import time +from typing import Any, Iterator, Optional, Union +from urllib.parse import urlencode + +from google.genai import _api_module +from google.genai import _common +from google.genai import types as genai_types +from google.genai._api_client import BaseApiClient +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 Pager + +from . import types + +logger = logging.getLogger("vertexai_genai.agentengines") + + +def _ReasoningEngineSpec_to_vertex( + api_client: BaseApiClient, + 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, ["agent_framework"]) is not None: + setv( + to_object, + ["agentFramework"], + getv(from_object, ["agent_framework"]), + ) + + if getv(from_object, ["class_methods"]) is not None: + setv(to_object, ["classMethods"], getv(from_object, ["class_methods"])) + + if getv(from_object, ["deployment_spec"]) is not None: + setv( + to_object, + ["deploymentSpec"], + getv(from_object, ["deployment_spec"]), + ) + + if getv(from_object, ["package_spec"]) is not None: + setv(to_object, ["packageSpec"], getv(from_object, ["package_spec"])) + + return to_object + + +def _CreateAgentEngineConfig_to_vertex( + api_client: BaseApiClient, + 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, ["spec"]) is not None: + setv( + parent_object, + ["spec"], + _ReasoningEngineSpec_to_vertex( + api_client, getv(from_object, ["spec"]), to_object + ), + ) + + return to_object + + +def _CreateAgentEngineRequestParameters_to_vertex( + api_client: BaseApiClient, + 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"], + _CreateAgentEngineConfig_to_vertex( + api_client, getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _GetOperationParameters_to_vertex( + api_client: BaseApiClient, + 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"]), + ) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + +def _DeleteAgentEngineConfig_to_vertex( + api_client: BaseApiClient, + 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, ["force"]) is not None: + setv(to_object, ["force"], getv(from_object, ["force"])) + + return to_object + + +def _DeleteAgentEngineRequestParameters_to_vertex( + api_client: BaseApiClient, + 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"], + _DeleteAgentEngineConfig_to_vertex( + api_client, getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _GetAgentEngineRequestParameters_to_vertex( + api_client: BaseApiClient, + 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 _ListAgentEngineConfig_to_vertex( + api_client: BaseApiClient, + 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"]), + ) + + if getv(from_object, ["filter"]) is not None: + setv(parent_object, ["_query", "filter"], getv(from_object, ["filter"])) + + return to_object + + +def _ListAgentEngineRequestParameters_to_vertex( + api_client: BaseApiClient, + 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"], + _ListAgentEngineConfig_to_vertex( + api_client, getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _QueryAgentEngineConfig_to_vertex( + api_client: BaseApiClient, + 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, ["class_method"]) is not None: + setv(parent_object, ["classMethod"], getv(from_object, ["class_method"])) + + if getv(from_object, ["input"]) is not None: + setv(parent_object, ["input"], getv(from_object, ["input"])) + + if getv(from_object, ["include_all_fields"]) is not None: + setv( + to_object, + ["includeAllFields"], + getv(from_object, ["include_all_fields"]), + ) + + return to_object + + +def _QueryAgentEngineRequestParameters_to_vertex( + api_client: BaseApiClient, + 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"], + _QueryAgentEngineConfig_to_vertex( + api_client, getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _UpdateAgentEngineConfig_to_vertex( + api_client: BaseApiClient, + 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, ["spec"]) is not None: + setv( + parent_object, + ["spec"], + _ReasoningEngineSpec_to_vertex( + api_client, getv(from_object, ["spec"]), to_object + ), + ) + + return to_object + + +def _UpdateAgentEngineRequestParameters_to_vertex( + api_client: BaseApiClient, + 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"], + _UpdateAgentEngineConfig_to_vertex( + api_client, getv(from_object, ["config"]), to_object + ), + ) + + return to_object + + +def _ReasoningEngine_from_vertex( + api_client: BaseApiClient, + 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, ["createTime"]) is not None: + setv(to_object, ["create_time"], getv(from_object, ["createTime"])) + + if getv(from_object, ["description"]) is not None: + setv(to_object, ["description"], getv(from_object, ["description"])) + + if getv(from_object, ["displayName"]) is not None: + setv(to_object, ["display_name"], getv(from_object, ["displayName"])) + + if getv(from_object, ["etag"]) is not None: + setv(to_object, ["etag"], getv(from_object, ["etag"])) + + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["spec"]) is not None: + setv(to_object, ["spec"], getv(from_object, ["spec"])) + + if getv(from_object, ["updateTime"]) is not None: + setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) + + return to_object + + +def _CreateAgentEngineOperation_from_vertex( + api_client: BaseApiClient, + 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, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + if getv(from_object, ["response"]) is not None: + setv( + to_object, + ["response"], + _ReasoningEngine_from_vertex( + api_client, getv(from_object, ["response"]), to_object + ), + ) + + return to_object + + +def _DeleteAgentEngineOperation_from_vertex( + api_client: BaseApiClient, + 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, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + return to_object + + +def _ListReasoningEnginesResponse_from_vertex( + api_client: BaseApiClient, + 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, ["nextPageToken"]) is not None: + setv(to_object, ["next_page_token"], getv(from_object, ["nextPageToken"])) + + if getv(from_object, ["reasoningEngines"]) is not None: + setv( + to_object, + ["reasoning_engines"], + [ + _ReasoningEngine_from_vertex(api_client, item, to_object) + for item in getv(from_object, ["reasoningEngines"]) + ], + ) + + return to_object + + +def _QueryReasoningEngineResponse_from_vertex( + api_client: BaseApiClient, + 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, ["output"]) is not None: + setv(to_object, ["output"], getv(from_object, ["output"])) + + return to_object + + +def _UpdateAgentEngineOperation_from_vertex( + api_client: BaseApiClient, + 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, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["metadata"]) is not None: + setv(to_object, ["metadata"], getv(from_object, ["metadata"])) + + if getv(from_object, ["done"]) is not None: + setv(to_object, ["done"], getv(from_object, ["done"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + return to_object + + +class AgentEngines(_api_module.BaseModule): + def _create( + self, *, config: Optional[types.CreateAgentEngineConfigOrDict] = None + ) -> types.CreateAgentEngineOperation: + """Creates a new Agent Engine.""" + + parameter_model = types._CreateAgentEngineRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _CreateAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + 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[genai_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_dict = self._api_client.request( + "post", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _CreateAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.CreateAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def _get_create_operation( + self, + *, + operation_name: str, + config: Optional[types.GetOperationConfigOrDict] = None, + ) -> types.CreateAgentEngineOperation: + parameter_model = types._GetOperationParameters( + 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 the Vertex AI client.") + else: + request_dict = _GetOperationParameters_to_vertex( + self._api_client, 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[genai_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_dict = self._api_client.request( + "get", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _CreateAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.CreateAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def delete( + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineConfigOrDict] = None, + ) -> types.DeleteAgentEngineOperation: + """Deletes an Agent Engine.""" + + parameter_model = types._DeleteAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _DeleteAgentEngineRequestParameters_to_vertex( + self._api_client, 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[genai_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_dict = self._api_client.request( + "delete", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _DeleteAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.DeleteAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def _get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineConfigOrDict] = None, + ) -> types.ReasoningEngine: + """Get an Agent Engine instance.""" + + parameter_model = types._GetAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _GetAgentEngineRequestParameters_to_vertex( + self._api_client, 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[genai_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_dict = self._api_client.request( + "get", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex( + self._api_client, response_dict + ) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def _list( + self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """Lists Agent Engines.""" + + parameter_model = types._ListAgentEngineRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _ListAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + 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[genai_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_dict = self._api_client.request( + "get", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex( + self._api_client, response_dict + ) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def _query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineConfigOrDict] = None, + ) -> types.QueryReasoningEngineResponse: + """Lists Agent Engines.""" + + parameter_model = types._QueryAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _QueryAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + 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[genai_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_dict = self._api_client.request( + "post", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _QueryReasoningEngineResponse_from_vertex( + self._api_client, response_dict + ) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def _update( + self, + *, + name: str, + config: Optional[types.UpdateAgentEngineConfigOrDict] = None, + ) -> types.UpdateAgentEngineOperation: + """Updates an Agent Engine.""" + + parameter_model = types._UpdateAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _UpdateAgentEngineRequestParameters_to_vertex( + self._api_client, 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[genai_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_dict = self._api_client.request( + "patch", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _UpdateAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.UpdateAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + def _list_pager( + self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None + ) -> Pager[types.ReasoningEngine]: + return Pager( + "reasoning_engines", + self._list, + self._list(config=config), + config, + ) + + def get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineConfigOrDict] = None, + ) -> types.AgentEngine: + """Gets an agent engine. + + Args: + name (str): Required. A fully-qualified resource name or ID such as + "projects/123/locations/us-central1/reasoningEngines/456" or a + shortened name such as "reasoningEngines/456". + """ + from vertexai.agent_engines import _agent_engines + + agent = types.AgentEngine( + api_client=self, + api_resource=self._get(name=name, config=config), + ) + try: + _agent_engines._register_api_methods_or_raise(agent) + except Exception as e: + logger.warning( + _agent_engines._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, + e, + ) + return agent + + def create( + self, + agent_engine: Any = None, + *, + staging_bucket: str, + requirements: Optional[Union[str, list[str]]] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + gcs_dir_name: Optional[str] = None, + extra_packages: Optional[list[str]] = None, + env_vars: Optional[Union[list[str], dict[str, Union[str, Any]]]] = None, + return_response: bool = True, + ) -> Union[types.AgentEngine, types.CreateAgentEngineOperation]: + """Creates an agent engine. + + The Agent Engine will be an instance of the `agent_engine` that + was passed in, running remotely on Vertex AI. + + Sample ``src_dir`` contents (e.g. ``./user_src_dir``): + + .. code-block:: python + + user_src_dir/ + |-- main.py + |-- requirements.txt + |-- user_code/ + | |-- utils.py + | |-- ... + |-- ... + + To build an Agent Engine with the above files, run: + + .. code-block:: python + + client = vertexai.genai.client.Client( + project="your-project", + location="us-central1", + ) + remote_agent = client.agent_engines.create( + agent_engine=local_agent, + requirements=[ + # I.e. the PyPI dependencies listed in requirements.txt + "google-cloud-aiplatform[agent_engines,adk]", + ... + ], + extra_packages=[ + "./user_src_dir/main.py", # a single file + "./user_src_dir/user_code", # a directory + ... + ], + ) + + Args: + staging_bucket (str): Required. The GCS bucket to use for staging + the artifacts needed. It must be a valid GCS bucket name, e.g. + "gs://bucket-name". + agent_engine (Any): Optional. The Agent Engine to be created. If not + specified, this will correspond to a lightweight instance that + cannot be queried (but can be updated to future instances that can + be queried). + requirements (Union[str, Sequence[str]]): Optional. The set of PyPI + dependencies needed. It can either be the path to a single file + (requirements.txt), or an ordered list of strings corresponding to + each line of the requirements file. + display_name (str): Optional. The user-defined name of the Agent + Engine. The name can be up to 128 characters long and can comprise + any UTF-8 character. + description (str): Optional. The description of the Agent Engine. + gcs_dir_name (str): Optional. The GCS bucket directory under + `staging_bucket` to use for staging the artifacts needed. + extra_packages (Sequence[str]): Optional. The set of extra + user-provided packages (if any). + env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]): + Optional. The environment variables to be set when running the + Agent Engine. If it is a list of strings, each string should be a + valid key to `os.environ`. If it is a dictionary, the keys are the + environment variable names, and the values are the corresponding + values. + return_response (bool): Optional. If True, the response will be + returned. Otherwise, the operation will be returned. + + Returns: + Union[types.AgentEngine, types.CreateAgentEngineOperation]: + It returns the Agent Engine if `return_response` is True, + otherwise + it returns the operation for creating the Agent Engine. + + Raises: + ValueError: If the `project` was not set using `client.Client`. + ValueError: If the `location` was not set using `client.Client`. + ValueError: If the `staging_bucket` was not set. + ValueError: If the `staging_bucket` does not start with "gs://". + ValueError: If `extra_packages` is specified but `agent_engine` is + None. + ValueError: If `requirements` is specified but `agent_engine` is + None. + ValueError: If `env_vars` has a dictionary entry that does not + correspond to a SecretRef. + ValueError: If `env_vars` is a list which contains a string that + does not exist in `os.environ`. + TypeError: If `env_vars` is not a list of strings or a dictionary. + TypeError: If `env_vars` has a value that is not a string or + SecretRef. + FileNotFoundError: If `extra_packages` includes a file or directory + that does not exist. + IOError: If requirements is a string that corresponds to a + nonexistent file. + """ + config = self._create_config( + agent_engine=agent_engine, + staging_bucket=staging_bucket, + requirements=requirements, + display_name=display_name, + description=description, + gcs_dir_name=gcs_dir_name, + extra_packages=extra_packages, + env_vars=env_vars, + ) + operation = self._create(config=config) + if return_response: + return self._await_create_operation(operation_name=operation.name) + return operation + + def _create_config( + self, + *, + agent_engine: Any = None, + staging_bucket: str, + requirements: Optional[Union[str, list[str]]] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + gcs_dir_name: Optional[str] = None, + extra_packages: Optional[list[str]] = None, + env_vars: Optional[Union[list[str], dict[str, Union[str, Any]]]] = None, + ): + import sys + from vertexai.agent_engines import _agent_engines + from vertexai.agent_engines import _utils + + sys_version = f"{sys.version_info.major}.{sys.version_info.minor}" + gcs_dir_name = gcs_dir_name or _agent_engines._DEFAULT_GCS_DIR_NAME + if agent_engine is not None: + agent_engine = _agent_engines._validate_agent_engine_or_raise(agent_engine) + _agent_engines._validate_staging_bucket_or_raise(staging_bucket) + if agent_engine is None: + if requirements is not None: + raise ValueError("requirements must be None if agent_engine is None.") + if extra_packages is not None: + raise ValueError("extra_packages must be None if agent_engine is None.") + requirements = _agent_engines._validate_requirements_or_raise( + agent_engine=agent_engine, + requirements=requirements, + ) + extra_packages = _agent_engines._validate_extra_packages_or_raise( + extra_packages + ) + # Prepares the Agent Engine for creation in Vertex AI. This involves + # packaging and uploading the artifacts for agent_engine, requirements and + # extra_packages to `staging_bucket/gcs_dir_name`. + _agent_engines._prepare( + agent_engine=agent_engine, + requirements=requirements, + project=self._api_client.project, + location=self._api_client.location, + staging_bucket=staging_bucket, + gcs_dir_name=gcs_dir_name, + extra_packages=extra_packages, + ) + config = {"display_name": display_name, "description": description} + if agent_engine is not None: + # Update the package spec. + package_spec = { + "python_version": sys_version, + "pickle_object_gcs_uri": "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _agent_engines._BLOB_FILENAME, + ), + } + if extra_packages: + package_spec["dependency_files_gcs_uri"] = "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _agent_engines._EXTRA_PACKAGES_FILE, + ) + if requirements: + package_spec["requirements_gcs_uri"] = "{}/{}/{}".format( + staging_bucket, + gcs_dir_name, + _agent_engines._REQUIREMENTS_FILE, + ) + agent_engine_spec = {"package_spec": package_spec} + if env_vars: + deployment_spec, _ = _agent_engines._generate_deployment_spec_or_raise( + env_vars=env_vars + ) + agent_engine_spec["deployment_spec"] = deployment_spec + class_methods = _agent_engines._generate_class_methods_spec_or_raise( + agent_engine=agent_engine, + operations=_agent_engines._get_registered_operations(agent_engine), + ) + agent_engine_spec["class_methods"] = [ + _utils.to_dict(class_method) for class_method in class_methods + ] + agent_engine_spec["agent_framework"] = _agent_engines._get_agent_framework( + agent_engine + ) + config["spec"] = agent_engine_spec + return config + + def _await_create_operation( + self, + *, + operation_name: str, + poll_interval_seconds: int = 10, + ) -> types.AgentEngine: + """Waits for the operation for creating an agent engine to complete. + + Args: + operation_name (str): Required. The name of the operation for + creating the Agent Engine. + poll_interval_seconds (int): The number of seconds to wait between + each poll. + + Returns: + AgentEngine: The Agent Engine that was created. + """ + from vertexai.agent_engines import _agent_engines + + operation = self._get_create_operation(operation_name=operation_name) + while not operation.done: + time.sleep(poll_interval_seconds) + operation = self._get_create_operation(operation_name=operation.name) + + agent = types.AgentEngine(api_client=self, api_resource=operation.response) + + try: + _agent_engines._register_api_methods_or_raise(agent) + except Exception as e: + logger.warning( + _agent_engines._FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, + e, + ) + return agent + + def list( + self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None + ) -> Iterator[types.AgentEngine]: + """Lists agent engines that matches the configuration.""" + + for reasoning_engine in self._list_pager(config=config): + yield types.AgentEngine( + api_client=self, + api_resource=reasoning_engine, + ) + + def _stream_query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineConfigOrDict] = None, + ) -> Iterator[Any]: + """Streams the response of the agent engine.""" + parameter_model = types._QueryAgentEngineRequestParameters( + name=name, + config=config, + ) + request_dict = _QueryAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:streamQuery?alt=sse".format_map(request_url_dict) + else: + path = "{name}:streamQuery?alt=sse" + 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[genai_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) + for response in self._api_client.request_streamed( + "post", path, request_dict, http_options + ): + yield response + + +class AsyncAgentEngines(_api_module.BaseModule): + async def _create( + self, *, config: Optional[types.CreateAgentEngineConfigOrDict] = None + ) -> types.CreateAgentEngineOperation: + """Creates a new Agent Engine.""" + + parameter_model = types._CreateAgentEngineRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _CreateAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + 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[genai_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_dict = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _CreateAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.CreateAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + async def _get_create_operation( + self, + *, + operation_name: str, + config: Optional[types.GetOperationConfigOrDict] = None, + ) -> types.CreateAgentEngineOperation: + parameter_model = types._GetOperationParameters( + 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 the Vertex AI client.") + else: + request_dict = _GetOperationParameters_to_vertex( + self._api_client, 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[genai_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_dict = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _CreateAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.CreateAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + async def delete( + self, + *, + name: str, + config: Optional[types.DeleteAgentEngineConfigOrDict] = None, + ) -> types.DeleteAgentEngineOperation: + """Deletes an Agent Engine.""" + + parameter_model = types._DeleteAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _DeleteAgentEngineRequestParameters_to_vertex( + self._api_client, 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[genai_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_dict = await self._api_client.async_request( + "delete", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _DeleteAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.DeleteAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + async def _get( + self, + *, + name: str, + config: Optional[types.GetAgentEngineConfigOrDict] = None, + ) -> types.ReasoningEngine: + """Get an Agent Engine instance.""" + + parameter_model = types._GetAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _GetAgentEngineRequestParameters_to_vertex( + self._api_client, 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[genai_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_dict = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _ReasoningEngine_from_vertex( + self._api_client, response_dict + ) + + return_value = types.ReasoningEngine._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + async def _list( + self, *, config: Optional[types.ListAgentEngineConfigOrDict] = None + ) -> types.ListReasoningEnginesResponse: + """Lists Agent Engines.""" + + parameter_model = types._ListAgentEngineRequestParameters( + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _ListAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "reasoningEngines".format_map(request_url_dict) + else: + path = "reasoningEngines" + + 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[genai_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_dict = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _ListReasoningEnginesResponse_from_vertex( + self._api_client, response_dict + ) + + return_value = types.ListReasoningEnginesResponse._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + async def _query( + self, + *, + name: str, + config: Optional[types.QueryAgentEngineConfigOrDict] = None, + ) -> types.QueryReasoningEngineResponse: + """Lists Agent Engines.""" + + parameter_model = types._QueryAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _QueryAgentEngineRequestParameters_to_vertex( + self._api_client, parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}:query".format_map(request_url_dict) + else: + path = "{name}:query" + + 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[genai_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_dict = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _QueryReasoningEngineResponse_from_vertex( + self._api_client, response_dict + ) + + return_value = types.QueryReasoningEngineResponse._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value + + async def _update( + self, + *, + name: str, + config: Optional[types.UpdateAgentEngineConfigOrDict] = None, + ) -> types.UpdateAgentEngineOperation: + """Updates an Agent Engine.""" + + parameter_model = types._UpdateAgentEngineRequestParameters( + 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 the Vertex AI client.") + else: + request_dict = _UpdateAgentEngineRequestParameters_to_vertex( + self._api_client, 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[genai_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_dict = await self._api_client.async_request( + "patch", path, request_dict, http_options + ) + + if self._api_client.vertexai: + response_dict = _UpdateAgentEngineOperation_from_vertex( + self._api_client, response_dict + ) + + return_value = types.UpdateAgentEngineOperation._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + self._api_client._verify_response(return_value) + return return_value diff --git a/vertexai/_genai/client.py b/vertexai/_genai/client.py index 3102c653c8..87b07fc453 100644 --- a/vertexai/_genai/client.py +++ b/vertexai/_genai/client.py @@ -30,6 +30,7 @@ def __init__(self, api_client: client.Client): self._api_client = api_client self._aio = AsyncClient(self._api_client) self._evals = None + self._agent_engines = None @property @_common.experimental_warning( @@ -50,6 +51,23 @@ def evals(self): ) from e return self._evals.AsyncEvals(self._api_client) + @property + def agent_engines(self): + if self._agent_engines is None: + try: + # We need to lazy load the agent_engines module to handle the + # possibility of ImportError when dependencies are not installed. + self._agent_engines = importlib.import_module( + ".agent_engines", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'agent_engines' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._agent_engines.AsyncAgentEngines(self._api_client) class Client: """Client for the GenAI SDK. @@ -101,6 +119,7 @@ def __init__( http_options=http_options, ) self._evals = None + self._agent_engines = None @property @_common.experimental_warning( @@ -120,3 +139,21 @@ def evals(self): "google-cloud-aiplatform[evaluation]" ) from e return self._evals.Evals(self._api_client) + + @property + def agent_engines(self): + if self._agent_engines is None: + try: + # We need to lazy load the agent_engines module to handle the + # possibility of ImportError when dependencies are not installed. + self._agent_engines = importlib.import_module( + ".agent_engines", + __package__, + ) + except ImportError as e: + raise ImportError( + "The 'agent_engines' module requires 'additional packages'. " + "Please install them using pip install " + "google-cloud-aiplatform[agent_engines]" + ) from e + return self._agent_engines.AgentEngines(self._api_client) diff --git a/vertexai/_genai/types.py b/vertexai/_genai/types.py index 506267c1dc..2f5fb1cccf 100644 --- a/vertexai/_genai/types.py +++ b/vertexai/_genai/types.py @@ -21,7 +21,17 @@ import logging import re import typing -from typing import Any, Callable, ClassVar, Literal, Optional, Tuple, TypeVar, Union +from typing import ( + Any, + Callable, + ClassVar, + Dict, + Literal, + Optional, + Tuple, + TypeVar, + Union, +) from google.genai import _common from google.genai import types as genai_types from pydantic import ( @@ -2382,6 +2392,816 @@ class EvaluateDatasetOperationDict(TypedDict, total=False): ] +class EnvVar(_common.BaseModel): + """Represents an environment variable present in a Container or Python Module.""" + + name: Optional[str] = Field( + default=None, + description="""Required. Name of the environment variable. Must be a valid C identifier.""", + ) + value: Optional[str] = Field( + default=None, + description="""Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""", + ) + + +class EnvVarDict(TypedDict, total=False): + """Represents an environment variable present in a Container or Python Module.""" + + name: Optional[str] + """Required. Name of the environment variable. Must be a valid C identifier.""" + + value: Optional[str] + """Required. Variables that reference a $(VAR_NAME) are expanded using the previous defined environment variables in the container and any service environment variables. If a variable cannot be resolved, the reference in the input string will be unchanged. The $(VAR_NAME) syntax can be escaped with a double $$, ie: $$(VAR_NAME). Escaped references will never be expanded, regardless of whether the variable exists or not.""" + + +EnvVarOrDict = Union[EnvVar, EnvVarDict] + + +class SecretRef(_common.BaseModel): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + + secret: Optional[str] = Field( + default=None, + description="""Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""", + ) + version: Optional[str] = Field( + default=None, + description="""The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""", + ) + + +class SecretRefDict(TypedDict, total=False): + """Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + + secret: Optional[str] + """Required. The name of the secret in Cloud Secret Manager. Format: {secret_name}.""" + + version: Optional[str] + """The Cloud Secret Manager secret version. Can be 'latest' for the latest version, an integer for a specific version, or a version alias.""" + + +SecretRefOrDict = Union[SecretRef, SecretRefDict] + + +class SecretEnvVar(_common.BaseModel): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" + + name: Optional[str] = Field( + default=None, + description="""Required. Name of the secret environment variable.""", + ) + secret_ref: Optional[SecretRef] = Field( + default=None, + description="""Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""", + ) + + +class SecretEnvVarDict(TypedDict, total=False): + """Represents an environment variable where the value is a secret in Cloud Secret Manager.""" + + name: Optional[str] + """Required. Name of the secret environment variable.""" + + secret_ref: Optional[SecretRefDict] + """Required. Reference to a secret stored in the Cloud Secret Manager that will provide the value for this environment variable.""" + + +SecretEnvVarOrDict = Union[SecretEnvVar, SecretEnvVarDict] + + +class ReasoningEngineSpecDeploymentSpec(_common.BaseModel): + """The specification of a Reasoning Engine deployment.""" + + env: Optional[list[EnvVar]] = Field( + default=None, + description="""Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""", + ) + secret_env: Optional[list[SecretEnvVar]] = Field( + default=None, + description="""Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""", + ) + + +class ReasoningEngineSpecDeploymentSpecDict(TypedDict, total=False): + """The specification of a Reasoning Engine deployment.""" + + env: Optional[list[EnvVarDict]] + """Optional. Environment variables to be set with the Reasoning Engine deployment. The environment variables can be updated through the UpdateReasoningEngine API.""" + + secret_env: Optional[list[SecretEnvVarDict]] + """Optional. Environment variables where the value is a secret in Cloud Secret Manager. To use this feature, add 'Secret Manager Secret Accessor' role (roles/secretmanager.secretAccessor) to AI Platform Reasoning Engine Service Agent.""" + + +ReasoningEngineSpecDeploymentSpecOrDict = Union[ + ReasoningEngineSpecDeploymentSpec, ReasoningEngineSpecDeploymentSpecDict +] + + +class ReasoningEngineSpecPackageSpec(_common.BaseModel): + """User provided package spec like pickled object and package requirements.""" + + dependency_files_gcs_uri: Optional[str] = Field( + default=None, + description="""Optional. The Cloud Storage URI of the dependency files in tar.gz format.""", + ) + pickle_object_gcs_uri: Optional[str] = Field( + default=None, + description="""Optional. The Cloud Storage URI of the pickled python object.""", + ) + python_version: Optional[str] = Field( + default=None, + description="""Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11. If not specified, default value is 3.10.""", + ) + requirements_gcs_uri: Optional[str] = Field( + default=None, + description="""Optional. The Cloud Storage URI of the `requirements.txt` file""", + ) + + +class ReasoningEngineSpecPackageSpecDict(TypedDict, total=False): + """User provided package spec like pickled object and package requirements.""" + + dependency_files_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the dependency files in tar.gz format.""" + + pickle_object_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the pickled python object.""" + + python_version: Optional[str] + """Optional. The Python version. Currently support 3.8, 3.9, 3.10, 3.11. If not specified, default value is 3.10.""" + + requirements_gcs_uri: Optional[str] + """Optional. The Cloud Storage URI of the `requirements.txt` file""" + + +ReasoningEngineSpecPackageSpecOrDict = Union[ + ReasoningEngineSpecPackageSpec, ReasoningEngineSpecPackageSpecDict +] + + +class ReasoningEngineSpec(_common.BaseModel): + """The specification of a Reasoning Engine.""" + + agent_framework: Optional[str] = Field( + default=None, + description="""Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""", + ) + class_methods: Optional[list[dict[str, Any]]] = Field( + default=None, + description="""Optional. Declarations for object class methods in OpenAPI specification format.""", + ) + deployment_spec: Optional[ReasoningEngineSpecDeploymentSpec] = Field( + default=None, + description="""Optional. The specification of a Reasoning Engine deployment.""", + ) + package_spec: Optional[ReasoningEngineSpecPackageSpec] = Field( + default=None, + description="""Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes.""", + ) + + +class ReasoningEngineSpecDict(TypedDict, total=False): + """The specification of a Reasoning Engine.""" + + agent_framework: Optional[str] + """Optional. The OSS agent framework used to develop the agent. Currently supported values: "google-adk", "langchain", "langgraph", "ag2", "llama-index", "custom".""" + + class_methods: Optional[list[dict[str, Any]]] + """Optional. Declarations for object class methods in OpenAPI specification format.""" + + deployment_spec: Optional[ReasoningEngineSpecDeploymentSpecDict] + """Optional. The specification of a Reasoning Engine deployment.""" + + package_spec: Optional[ReasoningEngineSpecPackageSpecDict] + """Optional. User provided package spec of the ReasoningEngine. Ignored when users directly specify a deployment image through `deployment_spec.first_party_image_override`, but keeping the field_behavior to avoid introducing breaking changes.""" + + +ReasoningEngineSpecOrDict = Union[ReasoningEngineSpec, ReasoningEngineSpecDict] + + +class CreateAgentEngineConfig(_common.BaseModel): + """Config for create agent engine.""" + + http_options: Optional[HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent Engine. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, + ) + description: Optional[str] = Field( + default=None, description="""The description of the Agent Engine.""" + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, + description="""Optional. Configurations of the ReasoningEngine.""", + ) + + +class CreateAgentEngineConfigDict(TypedDict, total=False): + """Config for create agent engine.""" + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + display_name: Optional[str] + """The user-defined name of the Agent Engine. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ + + description: Optional[str] + """The description of the Agent Engine.""" + + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the ReasoningEngine.""" + + +CreateAgentEngineConfigOrDict = Union[ + CreateAgentEngineConfig, CreateAgentEngineConfigDict +] + + +class _CreateAgentEngineRequestParameters(_common.BaseModel): + """Parameters for creating agent engines.""" + + config: Optional[CreateAgentEngineConfig] = Field(default=None, description="""""") + + +class _CreateAgentEngineRequestParametersDict(TypedDict, total=False): + """Parameters for creating agent engines.""" + + config: Optional[CreateAgentEngineConfigDict] + """""" + + +_CreateAgentEngineRequestParametersOrDict = Union[ + _CreateAgentEngineRequestParameters, _CreateAgentEngineRequestParametersDict +] + + +class ReasoningEngine(_common.BaseModel): + """An agent engine.""" + + create_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this ReasoningEngine was created.""", + ) + description: Optional[str] = Field( + default=None, + description="""Optional. The description of the ReasoningEngine.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""Required. The display name of the ReasoningEngine.""", + ) + etag: Optional[str] = Field( + default=None, + description="""Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""", + ) + name: Optional[str] = Field( + default=None, + description="""Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""", + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, + description="""Optional. Configurations of the ReasoningEngine""", + ) + update_time: Optional[datetime.datetime] = Field( + default=None, + description="""Output only. Timestamp when this ReasoningEngine was most recently updated.""", + ) + + +class ReasoningEngineDict(TypedDict, total=False): + """An agent engine.""" + + create_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngine was created.""" + + description: Optional[str] + """Optional. The description of the ReasoningEngine.""" + + display_name: Optional[str] + """Required. The display name of the ReasoningEngine.""" + + etag: Optional[str] + """Optional. Used to perform consistent read-modify-write updates. If not set, a blind "overwrite" update happens.""" + + name: Optional[str] + """Identifier. The resource name of the ReasoningEngine. Format: `projects/{project}/locations/{location}/reasoningEngines/{reasoning_engine}`""" + + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the ReasoningEngine""" + + update_time: Optional[datetime.datetime] + """Output only. Timestamp when this ReasoningEngine was most recently updated.""" + + +ReasoningEngineOrDict = Union[ReasoningEngine, ReasoningEngineDict] + + +class CreateAgentEngineOperation(_common.BaseModel): + """Operation for creating agent engines.""" + + 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[ReasoningEngine] = Field( + default=None, description="""The created Agent Engine.""" + ) + + +class CreateAgentEngineOperationDict(TypedDict, total=False): + """Operation for creating agent engines.""" + + 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[ReasoningEngineDict] + """The created Agent Engine.""" + + +CreateAgentEngineOperationOrDict = Union[ + CreateAgentEngineOperation, CreateAgentEngineOperationDict +] + + +class GetOperationConfig(_common.BaseModel): + + http_options: Optional[HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetOperationConfigDict(TypedDict, total=False): + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + +GetOperationConfigOrDict = Union[GetOperationConfig, GetOperationConfigDict] + + +class _GetOperationParameters(_common.BaseModel): + """Parameters for the GET method.""" + + operation_name: Optional[str] = Field( + default=None, + description="""The server-assigned name for the operation.""", + ) + config: Optional[GetOperationConfig] = Field( + default=None, + description="""Used to override the default configuration.""", + ) + + +class _GetOperationParametersDict(TypedDict, total=False): + """Parameters for the GET method.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetOperationConfigDict] + """Used to override the default configuration.""" + + +_GetOperationParametersOrDict = Union[ + _GetOperationParameters, _GetOperationParametersDict +] + + +class DeleteAgentEngineConfig(_common.BaseModel): + """Config for deleting agent engine.""" + + http_options: Optional[HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + force: Optional[bool] = Field( + default=None, + description="""If set to true, any child resources will also be deleted.""", + ) + + +class DeleteAgentEngineConfigDict(TypedDict, total=False): + """Config for deleting agent engine.""" + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + force: Optional[bool] + """If set to true, any child resources will also be deleted.""" + + +DeleteAgentEngineConfigOrDict = Union[ + DeleteAgentEngineConfig, DeleteAgentEngineConfigDict +] + + +class _DeleteAgentEngineRequestParameters(_common.BaseModel): + """Parameters for deleting agent engines.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent engine.""" + ) + config: Optional[DeleteAgentEngineConfig] = Field(default=None, description="""""") + + +class _DeleteAgentEngineRequestParametersDict(TypedDict, total=False): + """Parameters for deleting agent engines.""" + + name: Optional[str] + """Name of the agent engine.""" + + config: Optional[DeleteAgentEngineConfigDict] + """""" + + +_DeleteAgentEngineRequestParametersOrDict = Union[ + _DeleteAgentEngineRequestParameters, _DeleteAgentEngineRequestParametersDict +] + + +class DeleteAgentEngineOperation(_common.BaseModel): + """Operation for deleting agent engines.""" + + 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.""", + ) + + +class DeleteAgentEngineOperationDict(TypedDict, total=False): + """Operation for deleting agent engines.""" + + 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.""" + + +DeleteAgentEngineOperationOrDict = Union[ + DeleteAgentEngineOperation, DeleteAgentEngineOperationDict +] + + +class GetAgentEngineConfig(_common.BaseModel): + """Config for create agent engine.""" + + http_options: Optional[HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetAgentEngineConfigDict(TypedDict, total=False): + """Config for create agent engine.""" + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + +GetAgentEngineConfigOrDict = Union[GetAgentEngineConfig, GetAgentEngineConfigDict] + + +class _GetAgentEngineRequestParameters(_common.BaseModel): + """Parameters for getting agent engines.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent engine.""" + ) + config: Optional[GetAgentEngineConfig] = Field(default=None, description="""""") + + +class _GetAgentEngineRequestParametersDict(TypedDict, total=False): + """Parameters for getting agent engines.""" + + name: Optional[str] + """Name of the agent engine.""" + + config: Optional[GetAgentEngineConfigDict] + """""" + + +_GetAgentEngineRequestParametersOrDict = Union[ + _GetAgentEngineRequestParameters, _GetAgentEngineRequestParametersDict +] + + +class ListAgentEngineConfig(_common.BaseModel): + """Config for listing agent engines.""" + + http_options: Optional[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="""""") + filter: Optional[str] = Field(default=None, description="""""") + + +class ListAgentEngineConfigDict(TypedDict, total=False): + """Config for listing agent engines.""" + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + page_size: Optional[int] + """""" + + page_token: Optional[str] + """""" + + filter: Optional[str] + """""" + + +ListAgentEngineConfigOrDict = Union[ListAgentEngineConfig, ListAgentEngineConfigDict] + + +class _ListAgentEngineRequestParameters(_common.BaseModel): + """Parameters for listing agent engines.""" + + config: Optional[ListAgentEngineConfig] = Field(default=None, description="""""") + + +class _ListAgentEngineRequestParametersDict(TypedDict, total=False): + """Parameters for listing agent engines.""" + + config: Optional[ListAgentEngineConfigDict] + """""" + + +_ListAgentEngineRequestParametersOrDict = Union[ + _ListAgentEngineRequestParameters, _ListAgentEngineRequestParametersDict +] + + +class ListReasoningEnginesResponse(_common.BaseModel): + """Response for listing agent engines.""" + + next_page_token: Optional[str] = Field(default=None, description="""""") + reasoning_engines: Optional[list[ReasoningEngine]] = Field( + default=None, + description="""List of agent engines. + """, + ) + + +class ListReasoningEnginesResponseDict(TypedDict, total=False): + """Response for listing agent engines.""" + + next_page_token: Optional[str] + """""" + + reasoning_engines: Optional[list[ReasoningEngineDict]] + """List of agent engines. + """ + + +ListReasoningEnginesResponseOrDict = Union[ + ListReasoningEnginesResponse, ListReasoningEnginesResponseDict +] + + +class QueryAgentEngineConfig(_common.BaseModel): + """Config for querying agent engines.""" + + http_options: Optional[HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + class_method: Optional[str] = Field( + default=None, description="""The class method to call.""" + ) + input: Optional[dict[str, Any]] = Field( + default=None, description="""The input to the class method.""" + ) + include_all_fields: Optional[bool] = Field(default=False, description="""""") + + +class QueryAgentEngineConfigDict(TypedDict, total=False): + """Config for querying agent engines.""" + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + class_method: Optional[str] + """The class method to call.""" + + input: Optional[dict[str, Any]] + """The input to the class method.""" + + include_all_fields: Optional[bool] + """""" + + +QueryAgentEngineConfigOrDict = Union[QueryAgentEngineConfig, QueryAgentEngineConfigDict] + + +class _QueryAgentEngineRequestParameters(_common.BaseModel): + """Parameters for querying agent engines.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent engine.""" + ) + config: Optional[QueryAgentEngineConfig] = Field(default=None, description="""""") + + +class _QueryAgentEngineRequestParametersDict(TypedDict, total=False): + """Parameters for querying agent engines.""" + + name: Optional[str] + """Name of the agent engine.""" + + config: Optional[QueryAgentEngineConfigDict] + """""" + + +_QueryAgentEngineRequestParametersOrDict = Union[ + _QueryAgentEngineRequestParameters, _QueryAgentEngineRequestParametersDict +] + + +class QueryReasoningEngineResponse(_common.BaseModel): + """The response for querying an agent engine.""" + + output: Optional[Any] = Field( + default=None, + description="""Response provided by users in JSON object format.""", + ) + + +class QueryReasoningEngineResponseDict(TypedDict, total=False): + """The response for querying an agent engine.""" + + output: Optional[Any] + """Response provided by users in JSON object format.""" + + +QueryReasoningEngineResponseOrDict = Union[ + QueryReasoningEngineResponse, QueryReasoningEngineResponseDict +] + + +class UpdateAgentEngineConfig(_common.BaseModel): + """Config for updating agent engine.""" + + http_options: Optional[HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field( + default=None, + description="""The user-defined name of the Agent Engine. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """, + ) + description: Optional[str] = Field( + default=None, description="""The description of the Agent Engine.""" + ) + spec: Optional[ReasoningEngineSpec] = Field( + default=None, + description="""Optional. Configurations of the ReasoningEngine.""", + ) + + +class UpdateAgentEngineConfigDict(TypedDict, total=False): + """Config for updating agent engine.""" + + http_options: Optional[HttpOptionsDict] + """Used to override HTTP request options.""" + + display_name: Optional[str] + """The user-defined name of the Agent Engine. + + The display name can be up to 128 characters long and can comprise any + UTF-8 characters. + """ + + description: Optional[str] + """The description of the Agent Engine.""" + + spec: Optional[ReasoningEngineSpecDict] + """Optional. Configurations of the ReasoningEngine.""" + + +UpdateAgentEngineConfigOrDict = Union[ + UpdateAgentEngineConfig, UpdateAgentEngineConfigDict +] + + +class _UpdateAgentEngineRequestParameters(_common.BaseModel): + """Parameters for updating agent engines.""" + + name: Optional[str] = Field( + default=None, description="""Name of the agent engine.""" + ) + config: Optional[UpdateAgentEngineConfig] = Field(default=None, description="""""") + + +class _UpdateAgentEngineRequestParametersDict(TypedDict, total=False): + """Parameters for updating agent engines.""" + + name: Optional[str] + """Name of the agent engine.""" + + config: Optional[UpdateAgentEngineConfigDict] + """""" + + +_UpdateAgentEngineRequestParametersOrDict = Union[ + _UpdateAgentEngineRequestParameters, _UpdateAgentEngineRequestParametersDict +] + + +class UpdateAgentEngineOperation(_common.BaseModel): + """Operation for creating agent engines.""" + + 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.""", + ) + + +class UpdateAgentEngineOperationDict(TypedDict, total=False): + """Operation for creating agent engines.""" + + 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.""" + + +UpdateAgentEngineOperationOrDict = Union[ + UpdateAgentEngineOperation, UpdateAgentEngineOperationDict +] + + class PromptTemplate(_common.BaseModel): """A prompt template for creating prompts with variables.""" @@ -3073,3 +3893,48 @@ class EvaluateMethodConfigDict(TypedDict, total=False): EvaluateMethodConfigOrDict = Union[EvaluateMethodConfig, EvaluateMethodConfigDict] + + +class AgentEngine(_common.BaseModel): + """An agent engine instance.""" + + api_client: Optional[Any] = Field( + default=None, description="""The underlying API client.""" + ) + api_resource: Optional[ReasoningEngine] = Field( + default=None, + description="""The underlying API resource (i.e. ReasoningEngine).""", + ) + + # Allows dynamic binding of methods based on the registered operations. + model_config = ConfigDict(extra="allow") + + def __repr__(self) -> str: + return f"AgentEngine(api_resource.name='{self.api_resource.name}')" + + def operation_schemas(self) -> list[Dict[str, Any]]: + """Returns the schemas of all registered operations for the agent.""" + if not isinstance(self.api_resource, ReasoningEngine): + raise ValueError("api_resource is not initialized.") + if not self.api_resource.spec: + raise ValueError("api_resource.spec is not initialized.") + return self.api_resource.spec.class_methods + + def delete(self): + """Deletes the agent engine.""" + if not isinstance(self.api_resource, ReasoningEngine): + raise ValueError("api_resource is not initialized.") + self.api_client.delete(name=self.api_resource.name, config={"force": True}) + + +class AgentEngineDict(TypedDict, total=False): + """An agent engine instance.""" + + api_client: Optional[Any] + """The underlying API client.""" + + api_resource: Optional[ReasoningEngineDict] + """The underlying API resource (i.e. ReasoningEngine).""" + + +AgentEngineOrDict = Union[AgentEngine, AgentEngineDict] diff --git a/vertexai/agent_engines/_agent_engines.py b/vertexai/agent_engines/_agent_engines.py index cd7f0e3592..978e9b349e 100644 --- a/vertexai/agent_engines/_agent_engines.py +++ b/vertexai/agent_engines/_agent_engines.py @@ -1340,7 +1340,12 @@ def _unregister_api_methods( delattr(obj, method_name) -def _register_api_methods_or_raise(obj: "AgentEngine"): +def _register_api_methods_or_raise( + obj: "AgentEngine", + wrap_operation_fn: Optional[ + dict[str, Callable[[str, str], Callable[..., Any]]] + ] = None, +): """Registers Agent Engine API methods based on operation schemas. This function iterates through operation schemas provided by the @@ -1351,6 +1356,8 @@ def _register_api_methods_or_raise(obj: "AgentEngine"): Args: obj: The AgentEngine object to augment with API methods. + wrap_operation_fn: A dictionary of API modes and method wrapping + functions. Raises: ValueError: If the API mode is not supported or if the operation schema @@ -1369,66 +1376,49 @@ def _register_api_methods_or_raise(obj: "AgentEngine"): f" contain a `{_METHOD_NAME_KEY_IN_SCHEMA}` field." ) method_name = operation_schema.get(_METHOD_NAME_KEY_IN_SCHEMA) - method_description = operation_schema.get("description") - - if api_mode == _STANDARD_API_MODE: - method_description = ( - method_description - or _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_METHOD_NAME, - return_type=_DEFAULT_METHOD_RETURN_TYPE, - ) - ) - method = _wrap_query_operation( - method_name=method_name, - doc=method_description, - ) - elif api_mode == _ASYNC_API_MODE: - method_description = ( - method_description - or _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_ASYNC_METHOD_NAME, - return_type=_DEFAULT_ASYNC_METHOD_RETURN_TYPE, - ) - ) - method = _wrap_async_query_operation( - method_name=method_name, - doc=method_description, - ) - elif api_mode == _STREAM_API_MODE: - method_description = ( - method_description - or _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_STREAM_METHOD_NAME, - return_type=_DEFAULT_STREAM_METHOD_RETURN_TYPE, - ) - ) - method = _wrap_stream_query_operation( - method_name=method_name, - doc=method_description, - ) - elif api_mode == _ASYNC_STREAM_API_MODE: - method_description = ( - method_description - or _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( - method_name=method_name, - default_method_name=_DEFAULT_ASYNC_STREAM_METHOD_NAME, - return_type=_DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE, - ) - ) - method = _wrap_async_stream_query_operation( + default_method_name_map = { + _STANDARD_API_MODE: _DEFAULT_METHOD_NAME, + _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_NAME, + _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_NAME, + _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_NAME, + } + default_method_return_type_map = { + _STANDARD_API_MODE: _DEFAULT_METHOD_RETURN_TYPE, + _ASYNC_API_MODE: _DEFAULT_ASYNC_METHOD_RETURN_TYPE, + _STREAM_API_MODE: _DEFAULT_STREAM_METHOD_RETURN_TYPE, + _ASYNC_STREAM_API_MODE: _DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE, + } + method_description = operation_schema.get( + "description", + _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format( method_name=method_name, - doc=method_description, + default_method_name=default_method_name_map.get( + api_mode, + _DEFAULT_METHOD_NAME + ), + return_type=default_method_return_type_map.get( + api_mode, + _DEFAULT_METHOD_RETURN_TYPE, + ), ) + ) + _wrap_operation_map = { + _STANDARD_API_MODE: _wrap_query_operation, + _ASYNC_API_MODE: _wrap_async_query_operation, + _STREAM_API_MODE: _wrap_stream_query_operation, + _ASYNC_STREAM_API_MODE: _wrap_async_stream_query_operation, + } + if isinstance(wrap_operation_fn, dict) and api_mode in wrap_operation_fn: + # Override the default function with user-specified function if it exists. + _wrap_operation = wrap_operation_fn[api_mode] + elif api_mode in _wrap_operation_map: + _wrap_operation = _wrap_operation_map[api_mode] else: raise ValueError( f"Unsupported api mode: `{api_mode}`," - f" Supported modes are: `{_STANDARD_API_MODE}`, `{_ASYNC_API_MODE}`," - f" `{_STREAM_API_MODE}` and `{_ASYNC_STREAM_API_MODE}`." + f" Supported modes are: `{_wrap_operation_map.keys()}`." ) + method = _wrap_operation(method_name=method_name, doc=method_description) # Binds the method to the object. setattr(obj, method_name, types.MethodType(method, obj))