_agent_engines.py

# -*- coding: utf-8 -*-
# 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.
#
import abc
import inspect
import io
import json
import logging
import os
import sys
import tarfile
import types
import typing
from typing import (
    Any,
    AsyncIterable,
    Callable,
    Coroutine,
    Dict,
    Iterable,
    List,
    Optional,
    Protocol,
    Sequence,
    Tuple,
    Union,
)

from google.api_core import exceptions
from google.cloud import storage
from google.cloud.aiplatform import base
from google.cloud.aiplatform import initializer
from google.cloud.aiplatform import utils as aip_utils
from google.cloud.aiplatform_v1 import types as aip_types
from google.cloud.aiplatform_v1.types import reasoning_engine_service
from vertexai.agent_engines import _utils
import httpx
import proto

from google.protobuf import field_mask_pb2


_LOGGER = _utils.LOGGER
_SUPPORTED_PYTHON_VERSIONS = ("3.9", "3.10", "3.11", "3.12", "3.13", "3.14")
_DEFAULT_GCS_DIR_NAME = "agent_engine"
_BLOB_FILENAME = "agent_engine.pkl"
_REQUIREMENTS_FILE = "requirements.txt"
_EXTRA_PACKAGES_FILE = "dependencies.tar.gz"
_STANDARD_API_MODE = ""
_ASYNC_API_MODE = "async"
_STREAM_API_MODE = "stream"
_ASYNC_STREAM_API_MODE = "async_stream"
_BIDI_STREAM_API_MODE = "bidi_stream"
_A2A_EXTENSION_MODE = "a2a_extension"
_A2A_AGENT_CARD = "a2a_agent_card"
_MODE_KEY_IN_SCHEMA = "api_mode"
_METHOD_NAME_KEY_IN_SCHEMA = "name"
_DEFAULT_METHOD_NAME = "query"
_DEFAULT_ASYNC_METHOD_NAME = "async_query"
_DEFAULT_STREAM_METHOD_NAME = "stream_query"
_DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query"
_DEFAULT_METHOD_RETURN_TYPE = "dict[str, Any]"
_DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any, Any, Any]"
_DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]"
_DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]"
_DEFAULT_METHOD_DOCSTRING_TEMPLATE = """
    Runs the Agent Engine to serve the user request.
    This will be based on the `.{method_name}(...)` of the python object that
    was passed in when creating the Agent Engine. The method will invoke the
    `{default_method_name}` API client of the python object.
    Args:
        **kwargs:
            Optional. The arguments of the `.{method_name}(...)` method.
    Returns:
        {return_type}: The response from serving the user request.
"""
_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE = (
    "Failed to register API methods. Please follow the guide to "
    "register the API methods: "
    "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#custom-methods. "
    "Error: {%s}"
)
_AGENT_FRAMEWORK_ATTR = "agent_framework"
_DEFAULT_AGENT_FRAMEWORK = "custom"
_BUILD_OPTIONS_INSTALLATION = "installation_scripts"
_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,
}


try:
    from google.adk.agents import BaseAgent

    ADKAgent = BaseAgent
except (ImportError, AttributeError):
    ADKAgent = None

try:
    from a2a.types import (
        AgentCard,
        TransportProtocol,
        Message,
        TaskIdParams,
        TaskQueryParams,
    )
    from a2a.client import ClientConfig, ClientFactory

    AgentCard = AgentCard
    TransportProtocol = TransportProtocol
    Message = Message
    ClientConfig = ClientConfig
    ClientFactory = ClientFactory
    TaskIdParams = TaskIdParams
    TaskQueryParams = TaskQueryParams
except (ImportError, AttributeError):
    AgentCard = None
    TransportProtocol = None
    Message = None
    ClientConfig = None
    ClientFactory = None
    TaskIdParams = None
    TaskQueryParams = None


@typing.runtime_checkable
class Queryable(Protocol):
    """Protocol for Agent Engines that can be queried."""

    @abc.abstractmethod
    def query(self, **kwargs) -> Any:
        """Runs the Agent Engine to serve the user query."""


@typing.runtime_checkable
class AsyncQueryable(Protocol):
    """Protocol for Agent Engines that can be queried asynchronously."""

    @abc.abstractmethod
    def async_query(self, **kwargs) -> Coroutine[Any, Any, Any]:
        """Runs the Agent Engine to serve the user query asynchronously."""


@typing.runtime_checkable
class AsyncStreamQueryable(Protocol):
    """Protocol for Agent Engines that can stream responses asynchronously."""

    @abc.abstractmethod
    async def async_stream_query(self, **kwargs) -> AsyncIterable[Any]:
        """Asynchronously stream responses to serve the user query."""


@typing.runtime_checkable
class StreamQueryable(Protocol):
    """Protocol for Agent Engines that can stream responses."""

    @abc.abstractmethod
    def stream_query(self, **kwargs) -> Iterable[Any]:
        """Stream responses to serve the user query."""


@typing.runtime_checkable
class BidiStreamQueryable(Protocol):
    """Protocol for Agent Engines that can stream requests and responses."""

    @abc.abstractmethod
    async def bidi_stream_query(self, **kwargs) -> AsyncIterable[Any]:
        """Asynchronously stream requests and responses to serve the user query."""


@typing.runtime_checkable
class Cloneable(Protocol):
    """Protocol for Agent Engines that can be cloned."""

    @abc.abstractmethod
    def clone(self) -> Any:
        """Return a clone of the object."""


@typing.runtime_checkable
class OperationRegistrable(Protocol):
    """Protocol for agents that have registered operations."""

    @abc.abstractmethod
    def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:
        """Register the user provided operations (modes and methods)."""


_AgentEngineInterface = Union[
    ADKAgent,
    AsyncQueryable,
    AsyncStreamQueryable,
    BidiStreamQueryable,
    OperationRegistrable,
    Queryable,
    StreamQueryable,
]


def _wrap_agent_operation(agent: Any, operation: str):
    """Wraps an agent operation into a method (works for all API modes)."""

    def _method(self, **kwargs):
        if not self._tmpl_attrs.get("agent"):
            self.set_up()
        return getattr(self._tmpl_attrs["agent"], operation)(**kwargs)

    _method.__name__ = operation
    _method.__doc__ = getattr(agent, operation).__doc__
    return _method


class ModuleAgent(Cloneable, OperationRegistrable):
    """Agent that is defined by a module and an agent name.

    This agent is instantiated by importing a module and instantiating an agent
    from that module. It also allows to register operations that are defined in
    the agent.
    """

    def __init__(
        self,
        *,
        module_name: str,
        agent_name: str,
        register_operations: Dict[str, Sequence[str]],
        sys_paths: Optional[Sequence[str]] = None,
        agent_framework: Optional[str] = None,
    ):
        """Initializes a module-based agent.

        Args:
            module_name (str):
                Required. The name of the module to import.
            agent_name (str):
                Required. The name of the agent in the module to instantiate.
            register_operations (Dict[str, Sequence[str]]):
                Required. A dictionary of API modes to a list of method names.
            sys_paths (Sequence[str]):
                Optional. The system paths to search for the module. It should
                be relative to the directory where the code will be running.
                I.e. it should correspond to the directory being passed to
                `extra_packages=...` in the create method. It will be appended
                to the system path in the sequence being specified here, and
                only be appended if it is not already in the system path.
        """
        self.agent_framework = agent_framework
        self._tmpl_attrs = {
            "module_name": module_name,
            "agent_name": agent_name,
            "register_operations": register_operations,
            "sys_paths": sys_paths,
        }

    def clone(self):
        """Return a clone of the agent."""
        return ModuleAgent(
            module_name=self._tmpl_attrs.get("module_name"),
            agent_name=self._tmpl_attrs.get("agent_name"),
            register_operations=self._tmpl_attrs.get("register_operations"),
            sys_paths=self._tmpl_attrs.get("sys_paths"),
            agent_framework=self.agent_framework,
        )

    def register_operations(self, **kwargs) -> Dict[str, Sequence[str]]:
        return self._tmpl_attrs.get("register_operations")

    def set_up(self) -> None:
        """Sets up the agent for execution of queries at runtime.

        It runs the code to import the agent from the module, and registers the
        operations of the agent.
        """
        if self._tmpl_attrs.get("sys_paths"):
            import sys

            for sys_path in self._tmpl_attrs.get("sys_paths"):
                abs_path = os.path.abspath(sys_path)
                if abs_path not in sys.path:
                    sys.path.append(abs_path)

        import importlib

        module = importlib.import_module(self._tmpl_attrs.get("module_name"))
        try:
            importlib.reload(module)
        except Exception as e:
            _LOGGER.warning(
                f"Failed to reload module {self._tmpl_attrs.get('module_name')}: {e}"
            )
        agent_name = self._tmpl_attrs.get("agent_name")
        try:
            agent = getattr(module, agent_name)
        except AttributeError as e:
            raise AttributeError(
                f"Agent {agent_name} not found in module "
                f"{self._tmpl_attrs.get('module_name')}"
            ) from e
        if not self.agent_framework:
            self.agent_framework = _get_agent_framework(agent)
        self._tmpl_attrs["agent"] = agent
        if hasattr(agent, "set_up"):
            agent.set_up()
        for operations in self.register_operations().values():
            for operation in operations:
                op = _wrap_agent_operation(agent, operation)
                setattr(self, operation, types.MethodType(op, self))


class AgentEngine(base.VertexAiResourceNounWithFutureManager):
    """Represents a Vertex AI Agent Engine resource."""

    client_class = aip_utils.AgentEngineClientWithOverride
    _resource_noun = "reasoning_engine"
    _getter_method = "get_reasoning_engine"
    _list_method = "list_reasoning_engines"
    _delete_method = "delete_reasoning_engine"
    _parse_resource_name_method = "parse_reasoning_engine_path"
    _format_resource_name_method = "reasoning_engine_path"

    def __init__(self, resource_name: str):
        """Retrieves an Agent Engine resource.

        Args:
            resource_name (str):
                Required. A fully-qualified resource name or ID such as
                "projects/123/locations/us-central1/reasoningEngines/456" or
                "456" when project and location are initialized or passed.
        """
        super().__init__(resource_name=resource_name)
        self.execution_api_client = initializer.global_config.create_client(
            client_class=aip_utils.AgentEngineExecutionClientWithOverride,
        )
        self.execution_async_client = initializer.global_config.create_client(
            client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride,
        )
        self._gca_resource = self._get_gca_resource(resource_name=resource_name)
        try:
            _register_api_methods_or_raise(self)
        except Exception as e:
            _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e)
        self._operation_schemas = None

    @property
    def resource_name(self) -> str:
        """Fully-qualified resource name."""
        return self._gca_resource.name

    @classmethod
    def create(
        cls,
        agent_engine: Optional[_AgentEngineInterface] = None,
        *,
        requirements: Optional[Union[str, Sequence[str]]] = None,
        display_name: Optional[str] = None,
        description: Optional[str] = None,
        gcs_dir_name: Optional[str] = None,
        extra_packages: Optional[Sequence[str]] = None,
        env_vars: Optional[
            Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
        ] = None,
        build_options: Optional[Dict[str, Sequence[str]]] = None,
        service_account: Optional[str] = None,
        psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None,
        min_instances: Optional[int] = None,
        max_instances: Optional[int] = None,
        resource_limits: Optional[Dict[str, str]] = None,
        container_concurrency: Optional[int] = None,
        encryption_spec: Optional[aip_types.EncryptionSpec] = None,
    ) -> "AgentEngine":
        """Creates a new 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
            |   |-- ...
            |-- installation_scripts/
            |   |-- install_package.sh
            |   |-- ...
            |-- ...

        To build an Agent Engine with the above files, run:

        .. code-block:: python

            remote_agent = agent_engines.create(
                agent_engine=local_agent,
                requirements=[
                    # I.e. the PyPI dependencies listed in requirements.txt
                    "google-cloud-aiplatform==1.25.0",
                    "langchain==0.0.242",
                    ...
                ],
                extra_packages=[
                    "./user_src_dir/main.py", # a single file
                    "./user_src_dir/user_code", # a directory
                    ...
                ],
                build_options={
                    "installation_scripts": [
                        "./user_src_dir/installation_scripts/install_package.sh",
                        ...
                    ],
                },
            )

        Args:
            agent_engine (AgentEngineInterface):
                Optional. The Agent Engine to be created.
            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.
            build_options (Dict[str, Sequence[str]]):
                Optional. The build options for the Agent Engine.
                The following keys are supported:
                - installation_scripts:
                    Optional. The paths to the installation scripts to be
                    executed in the Docker image.
                    The scripts must be located in the `installation_scripts`
                    subdirectory and the path must be added to `extra_packages`.
            service_account (str):
                Optional. The service account to be used for the Agent Engine.
                If not specified, the default reasoning engine service agent
                service account will be used.
            psc_interface_config (aip_types.PscInterfaceConfig):
                Optional. The Private Service Connect interface config for the
                Agent Engine.
            min_instances (int):
                Optional. The minimum number of instances to be running for the
                Agent Engine.
            max_instances (int):
                Optional. The maximum number of instances to be running for the
                Agent Engine.
            resource_limits (Dict[str, str]):
                Optional. The resource limits for the Agent Engine.
            container_concurrency (int):
                Optional. The container concurrency for the Agent Engine.
            encryption_spec (aip_types.EncryptionSpec):
                Optional. The Cloud KMS resource identifier of the customer
                managed encryption key used to protect the model. Has the
                form:
                `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`.
                The key needs to be in the same region as the model.

        Returns:
            AgentEngine: The Agent Engine that was created.

        Raises:
            ValueError: If the `project` was not set using `vertexai.init`.
            ValueError: If the `location` was not set using `vertexai.init`.
            ValueError: If the `staging_bucket` was not set using vertexai.init.
            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.
        """
        sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
        _validate_sys_version_or_raise(sys_version)
        gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME
        staging_bucket = initializer.global_config.staging_bucket

        if agent_engine is not None:
            agent_engine = _validate_agent_engine_or_raise(agent_engine)
            staging_bucket = _validate_staging_bucket_or_raise(staging_bucket)
            if _is_adk_agent(None, agent_engine):
                env_vars = _add_telemetry_enablement_env(env_vars=env_vars)

        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 = _validate_requirements_or_raise(
            agent_engine=agent_engine,
            requirements=requirements,
        )
        extra_packages = _validate_extra_packages_or_raise(
            extra_packages=extra_packages,
            build_options=build_options,
        )

        sdk_resource = cls.__new__(cls)
        base.VertexAiResourceNounWithFutureManager.__init__(sdk_resource)

        # 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`.
        _prepare(
            agent_engine=agent_engine,
            requirements=requirements,
            project=sdk_resource.project,
            location=sdk_resource.location,
            staging_bucket=staging_bucket,
            gcs_dir_name=gcs_dir_name,
            extra_packages=extra_packages,
        )
        reasoning_engine = aip_types.ReasoningEngine(
            display_name=display_name,
            description=description,
            encryption_spec=encryption_spec,
        )
        if agent_engine is not None:
            # Update the package spec.
            package_spec = aip_types.ReasoningEngineSpec.PackageSpec(
                python_version=sys_version,
                pickle_object_gcs_uri="{}/{}/{}".format(
                    staging_bucket,
                    gcs_dir_name,
                    _BLOB_FILENAME,
                ),
            )
            if extra_packages:
                package_spec.dependency_files_gcs_uri = "{}/{}/{}".format(
                    staging_bucket,
                    gcs_dir_name,
                    _EXTRA_PACKAGES_FILE,
                )
            if requirements:
                package_spec.requirements_gcs_uri = "{}/{}/{}".format(
                    staging_bucket,
                    gcs_dir_name,
                    _REQUIREMENTS_FILE,
                )
            agent_engine_spec = aip_types.ReasoningEngineSpec(
                package_spec=package_spec,
            )
            if (
                env_vars
                or psc_interface_config
                or min_instances is not None
                or max_instances is not None
                or resource_limits
                or container_concurrency is not None
            ):
                deployment_spec, _ = _generate_deployment_spec_or_raise(
                    env_vars=env_vars,
                    psc_interface_config=psc_interface_config,
                    min_instances=min_instances,
                    max_instances=max_instances,
                    resource_limits=resource_limits,
                    container_concurrency=container_concurrency,
                )
                agent_engine_spec.deployment_spec = deployment_spec
            class_methods_spec = _generate_class_methods_spec_or_raise(
                agent_engine=agent_engine,
                operations=_get_registered_operations(agent_engine),
            )
            agent_engine_spec.class_methods.extend(class_methods_spec)
            if service_account:
                agent_engine_spec.service_account = service_account
            reasoning_engine.spec = agent_engine_spec
            reasoning_engine.spec.agent_framework = _get_agent_framework(agent_engine)
        operation_future = sdk_resource.api_client.create_reasoning_engine(
            parent=initializer.global_config.common_location_path(
                project=sdk_resource.project, location=sdk_resource.location
            ),
            reasoning_engine=reasoning_engine,
        )
        _LOGGER.log_create_with_lro(cls, operation_future)
        _LOGGER.info(
            f"View progress and logs at https://console.cloud.google.com/logs/query?project={sdk_resource.project}"
        )
        created_resource = operation_future.result()
        _LOGGER.info(f"{cls.__name__} created. Resource name: {created_resource.name}")
        _LOGGER.info(f"To use this {cls.__name__} in another session:")
        _LOGGER.info(
            f"agent_engine = vertexai.agent_engines.get('{created_resource.name}')"
        )
        # We use `._get_gca_resource(...)` instead of `created_resource` to
        # fully instantiate the attributes of the agent engine.
        sdk_resource._gca_resource = sdk_resource._get_gca_resource(
            resource_name=created_resource.name
        )
        sdk_resource.execution_api_client = initializer.global_config.create_client(
            client_class=aip_utils.AgentEngineExecutionClientWithOverride,
            credentials=sdk_resource.credentials,
            location_override=sdk_resource.location,
        )
        sdk_resource.execution_async_client = initializer.global_config.create_client(
            client_class=aip_utils.AgentEngineExecutionAsyncClientWithOverride,
            credentials=sdk_resource.credentials,
            location_override=sdk_resource.location,
        )
        if agent_engine is not None:
            try:
                _register_api_methods_or_raise(sdk_resource)
            except Exception as e:
                _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e)
        sdk_resource._operation_schemas = None
        return sdk_resource

    def update(
        self,
        *,
        agent_engine: Optional[_AgentEngineInterface] = None,
        requirements: Optional[Union[str, Sequence[str]]] = None,
        display_name: Optional[str] = None,
        description: Optional[str] = None,
        gcs_dir_name: Optional[str] = None,
        extra_packages: Optional[Sequence[str]] = None,
        env_vars: Optional[
            Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
        ] = None,
        build_options: Optional[Dict[str, Sequence[str]]] = None,
        service_account: Optional[str] = None,
        psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None,
        min_instances: Optional[int] = None,
        max_instances: Optional[int] = None,
        resource_limits: Optional[Dict[str, str]] = None,
        container_concurrency: Optional[int] = None,
        encryption_spec: Optional[aip_types.EncryptionSpec] = None,
    ) -> "AgentEngine":
        """Updates an existing Agent Engine.

        This method updates the configuration of an existing Agent Engine
        running remotely, which is identified by its resource name.
        Unlike the `create` function which requires a `agent_engine` object,
        all arguments in this method are optional.
        This method allows you to modify individual aspects of the configuration
        by providing any of the optional arguments.

        Args:
            agent_engine (AgentEngineInterface):
                Optional. The instance to be used as the updated Agent Engine.
                If it is not specified, the existing instance will be used.
            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.
                If it is not specified, the existing requirements will be used.
                If it is set to an empty string or list, the existing
                requirements will be removed.
            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). If
                it is not specified, the existing extra packages will be used.
                If it is set to an empty list, the existing extra packages will
                be removed.
            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.
            build_options (Dict[str, Sequence[str]]):
                Optional. The build options for the Agent Engine.
                The following keys are supported:
                - installation_scripts:
                    Optional. The paths to the installation scripts to be
                    executed in the Docker image.
                    The scripts must be located in the `installation_scripts`
                    subdirectory and the path must be added to `extra_packages`.
            service_account (str):
                Optional. The service account to be used for the Agent Engine.
                If not specified, the default reasoning engine service agent
                service account will be used.
            psc_interface_config (aip_types.PscInterfaceConfig):
                Optional. The Private Service Connect interface config for the
                Agent Engine.
            min_instances (int):
                Optional. The minimum number of instances to be running for the
                Agent Engine.
            max_instances (int):
                Optional. The maximum number of instances to be running for the
                Agent Engine.
            resource_limits (Dict[str, str]):
                Optional. The resource limits for the Agent Engine.
            container_concurrency (int):
                Optional. The container concurrency for the Agent Engine.
            encryption_spec (aip_types.EncryptionSpec):
                Optional. The Cloud KMS resource identifier of the customer
                managed encryption key used to protect the model. Has the
                form:
                `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`.
                The key needs to be in the same region as the model.

        Returns:
            AgentEngine: The Agent Engine that was updated.

        Raises:
            ValueError: If the `staging_bucket` was not set using vertexai.init.
            ValueError: If the `staging_bucket` does not start with "gs://".
            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.
            ValueError: if none of `display_name`, `description`, `requirements`,
            `extra_packages`, `env_vars`, or `agent_engine` were specified.
            IOError: If requirements is a string that corresponds to a
            nonexistent file.
        """
        staging_bucket = initializer.global_config.staging_bucket
        staging_bucket = _validate_staging_bucket_or_raise(staging_bucket)
        historical_operation_schemas = self.operation_schemas()
        gcs_dir_name = gcs_dir_name or _DEFAULT_GCS_DIR_NAME

        # Validate the arguments.
        if not any(
            [
                agent_engine,
                requirements,
                extra_packages,
                display_name,
                description,
                env_vars,
                build_options,
                service_account,
                psc_interface_config,
                min_instances is not None,
                max_instances is not None,
                resource_limits,
                container_concurrency is not None,
                encryption_spec,
            ]
        ):
            raise ValueError(
                "At least one of `agent_engine`, `requirements`, "
                "`extra_packages`, `display_name`, `description`, "
                "`env_vars`, `build_options`, `service_account`, "
                "`psc_interface_config`, `min_instances`, `max_instances`, "
                "`resource_limits`, `container_concurrency`, or "
                "`encryption_spec` must be specified."
            )
        if requirements is not None:
            requirements = _validate_requirements_or_raise(
                agent_engine=agent_engine,
                requirements=requirements,
            )
        if extra_packages is not None:
            extra_packages = _validate_extra_packages_or_raise(
                extra_packages=extra_packages,
                build_options=build_options,
            )
        if agent_engine is not None:
            agent_engine = _validate_agent_engine_or_raise(agent_engine)

        if _is_adk_agent(self, agent_engine):
            env_vars = _add_telemetry_enablement_env(env_vars=env_vars)

        # Prepares the Agent Engine for update in Vertex AI. This involves
        # packaging and uploading the artifacts for agent_engine, requirements
        # and extra_packages to `staging_bucket/gcs_dir_name`.
        _prepare(
            agent_engine=agent_engine,
            requirements=requirements,
            project=self.project,
            location=self.location,
            staging_bucket=staging_bucket,
            gcs_dir_name=gcs_dir_name,
            extra_packages=extra_packages,
        )
        update_request = _generate_update_request_or_raise(
            resource_name=self.resource_name,
            staging_bucket=staging_bucket,
            gcs_dir_name=gcs_dir_name,
            agent_engine=agent_engine,
            requirements=requirements,
            extra_packages=extra_packages,
            display_name=display_name,
            description=description,
            env_vars=env_vars,
            service_account=service_account,
            psc_interface_config=psc_interface_config,
            min_instances=min_instances,
            max_instances=max_instances,
            resource_limits=resource_limits,
            container_concurrency=container_concurrency,
            encryption_spec=encryption_spec,
        )
        operation_future = self.api_client.update_reasoning_engine(
            request=update_request
        )
        _LOGGER.info(
            f"Update Agent Engine backing LRO: {operation_future.operation.name}"
        )
        created_resource = operation_future.result()
        _LOGGER.info(f"Agent Engine updated. Resource name: {created_resource.name}")
        self._operation_schemas = None
        self.execution_api_client = initializer.global_config.create_client(
            client_class=aip_utils.AgentEngineExecutionClientWithOverride,
        )
        # We use `._get_gca_resource(...)` instead of `created_resource` to
        # fully instantiate the attributes of the agent engine.
        self._gca_resource = self._get_gca_resource(resource_name=self.resource_name)

        if (
            agent_engine is None
            or historical_operation_schemas == self.operation_schemas()
        ):
            # The operations of the agent engine are unchanged, so we return it.
            return self

        # If the agent engine has changed and the historical operation
        # schemas are different from the current operation schemas, we need to
        # unregister the historical operation schemas and register the current
        # operation schemas.
        _unregister_api_methods(self, historical_operation_schemas)
        try:
            _register_api_methods_or_raise(self)
        except Exception as e:
            _LOGGER.warning(_FAILED_TO_REGISTER_API_METHODS_WARNING_TEMPLATE, e)
        return self

    def delete(
        self,
        *,
        force: bool = False,
        **kwargs,
    ) -> None:
        """Deletes the ReasoningEngine.

        Args:
            force (bool):
                Optional. If set to True, child resources will also be deleted.
                Otherwise, the request will fail with FAILED_PRECONDITION error
                when the Agent Engine has undeleted child resources. Defaults to
                False.
            **kwargs (dict[str, Any]):
                Optional. Additional keyword arguments to pass to the
                delete_reasoning_engine method.
        """
        kwargs = kwargs or {}
        operation_future = self.api_client.delete_reasoning_engine(
            request=aip_types.DeleteReasoningEngineRequest(
                name=self.resource_name,
                force=force,
                **kwargs,
            ),
        )
        _LOGGER.info(
            f"Delete Agent Engine backing LRO: {operation_future.operation.name}"
        )
        operation_future.result()
        _LOGGER.info(f"Agent Engine deleted. Resource name: {self.resource_name}")

    def operation_schemas(self) -> Sequence[_utils.JsonDict]:
        """Returns the (Open)API schemas for the Agent Engine."""
        spec = _utils.to_dict(self._gca_resource.spec)
        if not hasattr(self, "_operation_schemas") or self._operation_schemas is None:
            self._operation_schemas = spec.get("class_methods", [])
        return self._operation_schemas


def _validate_sys_version_or_raise(sys_version: str) -> None:
    """Tries to validate the python system version."""
    if sys_version not in _SUPPORTED_PYTHON_VERSIONS:
        raise ValueError(
            f"Unsupported python version: {sys_version}. AgentEngine "
            f"only supports {_SUPPORTED_PYTHON_VERSIONS} at the moment."
        )
    if sys_version != f"{sys.version_info.major}.{sys.version_info.minor}":
        _LOGGER.warning(
            f"{sys_version=} is inconsistent with {sys.version_info=}. "
            "This might result in issues with deployment, and should only "
            "be used as a workaround for advanced cases."
        )


def _validate_staging_bucket_or_raise(staging_bucket: Optional[str]) -> str:
    """Tries to validate the staging bucket."""
    if not staging_bucket:
        raise ValueError("Please provide a `staging_bucket` in `vertexai.init(...)`")
    if not staging_bucket.startswith("gs://"):
        raise ValueError(f"{staging_bucket=} must start with `gs://`")
    return staging_bucket


def _validate_agent_engine_or_raise(
    agent_engine: _AgentEngineInterface,
    logger: base.Logger = _LOGGER,
) -> _AgentEngineInterface:
    """Tries to validate the agent engine.

    The agent engine must have one of the following:
    * a callable method named `query`
    * a callable method named `stream_query`
    * a callable method named `async_stream_query`
    * a callable method named `bidi_stream_query`
    * a callable method named `register_operations`

    Args:
        agent_engine: The agent engine to be validated.
        logger: The logger to use for logging.

    Returns:
        The validated agent engine.

    Raises:
        TypeError: If `agent_engine` has no callable method named `query`,
        `stream_query` or `register_operations`.
        ValueError: If `agent_engine` has an invalid `query`, `stream_query` or
        `register_operations` signature.
    """
    try:
        from google.adk.agents import BaseAgent

        if isinstance(agent_engine, BaseAgent):
            logger.info("Deploying google.adk.agents.Agent as an application.")
            from vertexai import agent_engines

            agent_engine = agent_engines.AdkApp(agent=agent_engine)
    except Exception:
        pass
    is_queryable = isinstance(agent_engine, Queryable) and callable(agent_engine.query)
    is_async_queryable = isinstance(agent_engine, AsyncQueryable) and callable(
        agent_engine.async_query
    )
    is_stream_queryable = isinstance(agent_engine, StreamQueryable) and callable(
        agent_engine.stream_query
    )
    is_async_stream_queryable = isinstance(
        agent_engine, AsyncStreamQueryable
    ) and callable(agent_engine.async_stream_query)
    is_bidi_stream_queryable = isinstance(
        agent_engine, BidiStreamQueryable
    ) and callable(agent_engine.bidi_stream_query)
    is_operation_registrable = isinstance(
        agent_engine, OperationRegistrable
    ) and callable(agent_engine.register_operations)

    if not (
        is_queryable
        or is_async_queryable
        or is_stream_queryable
        or is_operation_registrable
        or is_async_stream_queryable
        or is_bidi_stream_queryable
    ):
        raise TypeError(
            "agent_engine has none of the following callable methods: "
            "`query`, `async_query`, `stream_query`, `async_stream_query`, "
            "`bidi_stream_query` or `register_operations`."
        )

    if is_queryable:
        try:
            inspect.signature(getattr(agent_engine, "query"))
        except ValueError as err:
            raise ValueError(
                "Invalid query signature. This might be due to a missing "
                "`self` argument in the agent_engine.query method."
            ) from err

    if is_async_queryable:
        try:
            inspect.signature(getattr(agent_engine, "async_query"))
        except ValueError as err:
            raise ValueError(
                "Invalid async_query signature. This might be due to a missing "
                "`self` argument in the agent_engine.async_query method."
            ) from err

    if is_stream_queryable:
        try:
            inspect.signature(getattr(agent_engine, "stream_query"))
        except ValueError as err:
            raise ValueError(
                "Invalid stream_query signature. This might be due to a missing"
                " `self` argument in the agent_engine.stream_query method."
            ) from err

    if is_async_stream_queryable:
        try:
            inspect.signature(getattr(agent_engine, "async_stream_query"))
        except ValueError as err:
            raise ValueError(
                "Invalid async_stream_query signature. This might be due to a "
                " missing `self` argument in the "
                "agent_engine.async_stream_query method."
            ) from err

    if is_bidi_stream_queryable:
        try:
            inspect.signature(getattr(agent_engine, "bidi_stream_query"))
        except ValueError as err:
            raise ValueError(
                "Invalid bidi_stream_query signature. This might be due to a "
                " missing `self` argument in the "
                "agent_engine.bidi_stream_query method."
            ) from err

    if is_operation_registrable:
        try:
            inspect.signature(getattr(agent_engine, "register_operations"))
        except ValueError as err:
            raise ValueError(
                "Invalid register_operations signature. This might be due to a "
                "missing `self` argument in the "
                "agent_engine.register_operations method."
            ) from err

    if isinstance(agent_engine, Cloneable):
        # Avoid undeployable states.
        agent_engine = agent_engine.clone()
    return agent_engine


def _is_adk_agent(
    agent_engine_to_update: Optional[AgentEngine],
    new_agent_engine: Optional[_AgentEngineInterface],
) -> bool:
    """Checks if the agent engine is an ADK agent.

    Args:
        agent_engine_to_update: Existing agent engine, None if creating new one.
        new_agent_engine: The new agent engine to deploy. Can be None during an update, if the Python agent implementation is not provided, and should remain unchanged.

    Returns:
        True if the agent after the create/update operation, will be an ADK agent.
    """

    from vertexai.agent_engines.templates import adk

    if new_agent_engine is not None:
        return (
            getattr(new_agent_engine, "agent_framework", None)
            == adk.AdkApp.agent_framework
        )
    if agent_engine_to_update is not None:
        return (
            agent_engine_to_update.gca_resource.spec.agent_framework
            == adk.AdkApp.agent_framework
        )
    return False


EnvVars = Optional[Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]]


def _add_telemetry_enablement_env(*, env_vars: EnvVars) -> EnvVars:
    """Adds telemetry enablement env var to the env vars.

    This is in order to achieve default-on telemetry.
    If the telemetry enablement env var is already set, we do not override it.

    Args:
        env_vars: The env vars to add the telemetry enablement env var to.

    Returns:
        The env vars with the telemetry enablement env var added.
    """

    GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY = (
        "GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY"
    )

    if env_vars is None:
        return {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}
    if isinstance(env_vars, dict):
        return (
            env_vars
            if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars
            else env_vars | {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}
        )
    if isinstance(env_vars, list) or isinstance(env_vars, tuple):
        if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY not in os.environ:
            os.environ[GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY] = "unspecified"

        if isinstance(env_vars, list):
            return env_vars + [GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY]
        else:
            return env_vars + (GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY,)

    raise TypeError(
        f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}."
    )


def _validate_requirements_or_raise(
    *,
    agent_engine: _AgentEngineInterface,
    requirements: Optional[Sequence[str]] = None,
    logger: logging.getLoggerClass() = _LOGGER,
) -> Sequence[str]:
    """Tries to validate the requirements."""
    if requirements is None:
        requirements = []
    elif isinstance(requirements, str):
        try:
            logger.info(f"Reading requirements from {requirements=}")
            with open(requirements) as f:
                requirements = f.read().splitlines()
                logger.info(f"Read the following lines: {requirements}")
        except IOError as err:
            raise IOError(f"Failed to read requirements from {requirements=}") from err
    requirements = _utils.validate_requirements_or_warn(
        obj=agent_engine,
        requirements=requirements,
        logger=logger,
    )
    logger.info(f"The final list of requirements: {requirements}")
    return requirements


def _validate_extra_packages_or_raise(
    extra_packages: Optional[Sequence[str]],
    build_options: Optional[Dict[str, Sequence[str]]] = None,
) -> Sequence[str]:
    """Tries to validates the extra packages."""
    extra_packages = extra_packages or []
    if build_options and _BUILD_OPTIONS_INSTALLATION in build_options:
        _utils.validate_installation_scripts_or_raise(
            script_paths=build_options[_BUILD_OPTIONS_INSTALLATION],
            extra_packages=extra_packages,
        )
    for extra_package in extra_packages:
        if not os.path.exists(extra_package):
            raise FileNotFoundError(
                f"Extra package specified but not found: {extra_package=}"
            )
    return extra_packages


def _get_gcs_bucket(
    *,
    project: str,
    location: str,
    staging_bucket: str,
    logger: base.Logger = _LOGGER,
) -> storage.Bucket:
    """Gets or creates the GCS bucket."""
    storage = _utils._import_cloud_storage_or_raise()
    storage_client = storage.Client(project=project)
    staging_bucket = staging_bucket.replace("gs://", "")
    try:
        gcs_bucket = storage_client.get_bucket(staging_bucket)
        logger.info(f"Using bucket {staging_bucket}")
    except exceptions.NotFound:
        new_bucket = storage_client.bucket(staging_bucket)
        gcs_bucket = storage_client.create_bucket(new_bucket, location=location)
        logger.info(f"Creating bucket {staging_bucket} in {location=}")
    return gcs_bucket


def _upload_agent_engine(
    *,
    agent_engine: _AgentEngineInterface,
    gcs_bucket: storage.Bucket,
    gcs_dir_name: str,
    logger: base.Logger = _LOGGER,
) -> None:
    """Uploads the agent engine to GCS."""
    cloudpickle = _utils._import_cloudpickle_or_raise()
    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}")
    with blob.open("wb") as f:
        try:
            cloudpickle.dump(agent_engine, f)
        except Exception as e:
            url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations"
            raise TypeError(
                f"Failed to serialize agent engine. Visit {url} for details."
            ) from e
    with blob.open("rb") as f:
        try:
            _ = cloudpickle.load(f)
        except Exception as e:
            raise TypeError("Agent engine serialized to an invalid format") from e
    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"
    logger.info(f"Wrote to {dir_name}/{_BLOB_FILENAME}")


def _upload_requirements(
    *,
    requirements: Sequence[str],
    gcs_bucket: storage.Bucket,
    gcs_dir_name: str,
    logger: base.Logger = _LOGGER,
) -> None:
    """Uploads the requirements file to GCS."""
    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_REQUIREMENTS_FILE}")
    blob.upload_from_string("\n".join(requirements))
    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"
    logger.info(f"Writing to {dir_name}/{_REQUIREMENTS_FILE}")


def _upload_extra_packages(
    *,
    extra_packages: Sequence[str],
    gcs_bucket: storage.Bucket,
    gcs_dir_name: str,
    logger: base.Logger = _LOGGER,
) -> None:
    """Uploads extra packages to GCS."""
    logger.info("Creating in-memory tarfile of extra_packages")
    tar_fileobj = io.BytesIO()
    with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar:
        for file in extra_packages:
            tar.add(file)
    tar_fileobj.seek(0)
    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_EXTRA_PACKAGES_FILE}")
    blob.upload_from_string(tar_fileobj.read())
    dir_name = f"gs://{gcs_bucket.name}/{gcs_dir_name}"
    logger.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}")


def _prepare(
    agent_engine: Optional[_AgentEngineInterface],
    requirements: Optional[Sequence[str]],
    extra_packages: Optional[Sequence[str]],
    project: str,
    location: str,
    staging_bucket: str,
    gcs_dir_name: str,
    logger: base.Logger = _LOGGER,
) -> None:
    """Prepares the agent engine for creation or updates in Vertex AI.

    This involves packaging and uploading artifacts to Cloud Storage. Note that
    1. This does not actually update the Agent Engine in Vertex AI.
    2. This will only generate and upload a pickled object if specified.
    3. This will only generate and upload the dependencies.tar.gz file if
    extra_packages is non-empty.

    Args:
        agent_engine: The agent engine to be prepared.
        requirements (Sequence[str]): The set of PyPI dependencies needed.
        extra_packages (Sequence[str]): The set of extra user-provided packages.
        project (str): The project for the staging bucket.
        location (str): The location for the staging bucket.
        staging_bucket (str): The staging bucket name in the form "gs://...".
        gcs_dir_name (str): The GCS bucket directory under `staging_bucket` to
            use for staging the artifacts needed.
    """
    if agent_engine is None:
        return
    gcs_bucket = _get_gcs_bucket(
        project=project,
        location=location,
        staging_bucket=staging_bucket,
        logger=logger,
    )
    _upload_agent_engine(
        agent_engine=agent_engine,
        gcs_bucket=gcs_bucket,
        gcs_dir_name=gcs_dir_name,
        logger=logger,
    )
    if requirements is not None:
        _upload_requirements(
            requirements=requirements,
            gcs_bucket=gcs_bucket,
            gcs_dir_name=gcs_dir_name,
            logger=logger,
        )
    if extra_packages is not None:
        _upload_extra_packages(
            extra_packages=extra_packages,
            gcs_bucket=gcs_bucket,
            gcs_dir_name=gcs_dir_name,
            logger=logger,
        )


def _update_deployment_spec_with_env_vars_dict_or_raise(
    *,
    deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec,
    env_vars: Dict[str, Union[str, aip_types.SecretRef]],
) -> None:
    for key, value in env_vars.items():
        if isinstance(value, Dict):
            try:
                secret_ref = _utils.to_proto(value, aip_types.SecretRef())
            except Exception as e:
                raise ValueError(f"Failed to convert to secret ref: {value}") from e
            deployment_spec.secret_env.append(
                aip_types.SecretEnvVar(name=key, secret_ref=secret_ref)
            )
        elif isinstance(value, aip_types.SecretRef):
            deployment_spec.secret_env.append(
                aip_types.SecretEnvVar(name=key, secret_ref=value)
            )
        elif isinstance(value, str):
            deployment_spec.env.append(aip_types.EnvVar(name=key, value=value))
        else:
            raise TypeError(
                f"Unknown value type in env_vars for {key}. "
                f"Must be a str or SecretRef: {value}"
            )


def _update_deployment_spec_with_env_vars_list_or_raise(
    *,
    deployment_spec: aip_types.ReasoningEngineSpec.DeploymentSpec,
    env_vars: Sequence[str],
) -> None:
    for env_var in env_vars:
        if env_var not in os.environ:
            raise ValueError(f"Env var not found in os.environ: {env_var}.")
        deployment_spec.env.append(
            aip_types.EnvVar(name=env_var, value=os.environ[env_var])
        )


def _generate_deployment_spec_or_raise(
    *,
    env_vars: Optional[
        Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
    ] = None,
    psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None,
    min_instances: Optional[int] = None,
    max_instances: Optional[int] = None,
    resource_limits: Optional[Dict[str, str]] = None,
    container_concurrency: Optional[int] = None,
) -> Tuple[aip_types.ReasoningEngineSpec.DeploymentSpec, List[str]]:
    deployment_spec = aip_types.ReasoningEngineSpec.DeploymentSpec()
    update_masks = []
    if env_vars:
        deployment_spec.env = []
        deployment_spec.secret_env = []
        if isinstance(env_vars, Dict):
            _update_deployment_spec_with_env_vars_dict_or_raise(
                deployment_spec=deployment_spec,
                env_vars=env_vars,
            )
        elif isinstance(env_vars, Sequence):
            _update_deployment_spec_with_env_vars_list_or_raise(
                deployment_spec=deployment_spec,
                env_vars=env_vars,
            )
        else:
            raise TypeError(
                f"env_vars must be a list, tuple or a dict, but got {type(env_vars)}."
            )
        if deployment_spec.env:
            update_masks.append("spec.deployment_spec.env")
        if deployment_spec.secret_env:
            update_masks.append("spec.deployment_spec.secret_env")
    if psc_interface_config:
        deployment_spec.psc_interface_config = psc_interface_config
        update_masks.append("spec.deployment_spec.psc_interface_config")
    if min_instances is not None:
        deployment_spec.min_instances = min_instances
        update_masks.append("spec.deployment_spec.min_instances")
    if max_instances is not None:
        deployment_spec.max_instances = max_instances
        update_masks.append("spec.deployment_spec.max_instances")
    if resource_limits:
        deployment_spec.resource_limits = resource_limits
        update_masks.append("spec.deployment_spec.resource_limits")
    if container_concurrency is not None:
        deployment_spec.container_concurrency = container_concurrency
        update_masks.append("spec.deployment_spec.container_concurrency")
    return deployment_spec, update_masks


def _get_agent_framework(
    agent_engine: _AgentEngineInterface,
) -> str:
    if (
        hasattr(agent_engine, _AGENT_FRAMEWORK_ATTR)
        and getattr(agent_engine, _AGENT_FRAMEWORK_ATTR) is not None
        and isinstance(getattr(agent_engine, _AGENT_FRAMEWORK_ATTR), str)
    ):
        return getattr(agent_engine, _AGENT_FRAMEWORK_ATTR)
    return _DEFAULT_AGENT_FRAMEWORK


def _generate_update_request_or_raise(
    *,
    resource_name: str,
    staging_bucket: str,
    gcs_dir_name: str = _DEFAULT_GCS_DIR_NAME,
    agent_engine: Optional[_AgentEngineInterface] = None,
    requirements: Optional[Union[str, Sequence[str]]] = None,
    extra_packages: Optional[Sequence[str]] = None,
    display_name: Optional[str] = None,
    description: Optional[str] = None,
    env_vars: Optional[
        Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
    ] = None,
    service_account: Optional[str] = None,
    psc_interface_config: Optional[aip_types.PscInterfaceConfig] = None,
    min_instances: Optional[int] = None,
    max_instances: Optional[int] = None,
    resource_limits: Optional[Dict[str, str]] = None,
    container_concurrency: Optional[int] = None,
    encryption_spec: Optional[str] = None,
) -> reasoning_engine_service.UpdateReasoningEngineRequest:
    """Tries to generate the update request for the agent engine."""
    is_spec_update = False
    update_masks: List[str] = []
    agent_engine_spec = aip_types.ReasoningEngineSpec()
    package_spec = aip_types.ReasoningEngineSpec.PackageSpec()
    if requirements is not None:
        is_spec_update = True
        update_masks.append("spec.package_spec.requirements_gcs_uri")
        package_spec.requirements_gcs_uri = "{}/{}/{}".format(
            staging_bucket,
            gcs_dir_name,
            _REQUIREMENTS_FILE,
        )
    if extra_packages is not None:
        is_spec_update = True
        update_masks.append("spec.package_spec.dependency_files_gcs_uri")
        package_spec.dependency_files_gcs_uri = "{}/{}/{}".format(
            staging_bucket,
            gcs_dir_name,
            _EXTRA_PACKAGES_FILE,
        )
    if agent_engine is not None:
        is_spec_update = True
        update_masks.append("spec.package_spec.pickle_object_gcs_uri")
        package_spec.pickle_object_gcs_uri = "{}/{}/{}".format(
            staging_bucket,
            gcs_dir_name,
            _BLOB_FILENAME,
        )
        class_methods_spec = _generate_class_methods_spec_or_raise(
            agent_engine=agent_engine,
            operations=_get_registered_operations(agent_engine),
        )
        agent_engine_spec.class_methods.extend(class_methods_spec)
        update_masks.append("spec.class_methods")
        agent_engine_spec.agent_framework = _get_agent_framework(agent_engine)
        update_masks.append("spec.agent_framework")
    if (
        env_vars is not None
        or psc_interface_config
        or min_instances is not None
        or max_instances is not None
        or resource_limits
        or container_concurrency is not None
    ):
        is_spec_update = True
        deployment_spec, deployment_update_masks = _generate_deployment_spec_or_raise(
            env_vars=env_vars,
            psc_interface_config=psc_interface_config,
            min_instances=min_instances,
            max_instances=max_instances,
            resource_limits=resource_limits,
            container_concurrency=container_concurrency,
        )
        update_masks.extend(deployment_update_masks)
        agent_engine_spec.deployment_spec = deployment_spec
    if service_account is not None:
        is_spec_update = True
        update_masks.append("spec.service_account")
        agent_engine_spec.service_account = service_account

    agent_engine_message = aip_types.ReasoningEngine(name=resource_name)
    if is_spec_update:
        if package_spec:
            agent_engine_spec.package_spec = package_spec
        agent_engine_message.spec = agent_engine_spec
    if display_name:
        agent_engine_message.display_name = display_name
        update_masks.append("display_name")
    if description:
        agent_engine_message.description = description
        update_masks.append("description")
    if encryption_spec:
        agent_engine_message.encryption_spec = encryption_spec
        update_masks.append("encryption_spec")
    if not update_masks:
        raise ValueError(
            "At least one of `agent_engine`, `requirements`, `extra_packages`, "
            "`display_name`, `description`, `env_vars`, or "
            "`encryption_spec` must be specified."
        )
    return reasoning_engine_service.UpdateReasoningEngineRequest(
        reasoning_engine=agent_engine_message,
        update_mask=field_mask_pb2.FieldMask(paths=update_masks),
    )


def _wrap_query_operation(method_name: str) -> Callable[..., _utils.JsonDict]:
    """Wraps an Agent Engine method, creating a callable for `query` API.

    This function creates a callable object that executes the specified
    Agent Engine method using the `query` API.  It handles the creation of
    the API request and the processing of the API response.

    Args:
        method_name: The name of the Agent Engine method to call.
        doc: Documentation string for the method.

    Returns:
        A callable object that executes the method on the Agent Engine via
        the `query` API.
    """

    def _method(self, **kwargs) -> _utils.JsonDict:
        response = self.execution_api_client.query_reasoning_engine(
            request=aip_types.QueryReasoningEngineRequest(
                name=self.resource_name,
                input=kwargs,
                class_method=method_name,
            ),
        )
        output = _utils.to_dict(response)
        return output.get("output", output)

    return _method


def _wrap_async_query_operation(method_name: str) -> Callable[..., Coroutine]:
    """Wraps an Agent Engine method, creating an async callable for `query` API.

    This function creates a callable object that executes the specified
    Agent Engine method asynchronously using the `query` API. It handles the
    creation of the API request and the processing of the API response.

    Args:
        method_name: The name of the Agent Engine method to call.
        doc: Documentation string for the method.

    Returns:
        A callable object that executes the method on the Agent Engine via
        the `query` API.
    """

    async def _method(self, **kwargs) -> _utils.JsonDict:
        response = await self.execution_async_client.query_reasoning_engine(
            request=aip_types.QueryReasoningEngineRequest(
                name=self.resource_name,
                input=kwargs,
                class_method=method_name,
            ),
        )
        output = _utils.to_dict(response)
        return output.get("output", output)

    return _method


def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterable[Any]]:
    """Wraps an Agent Engine method, creating a callable for `stream_query` API.

    This function creates a callable object that executes the specified
    Agent Engine method using the `stream_query` API.  It handles the
    creation of the API request and the processing of the API response.

    Args:
        method_name: The name of the Agent Engine method to call.
        doc: Documentation string for the method.

    Returns:
        A callable object that executes the method on the Agent Engine via
        the `stream_query` API.
    """

    def _method(self, **kwargs) -> Iterable[Any]:
        response = self.execution_api_client.stream_query_reasoning_engine(
            request=aip_types.StreamQueryReasoningEngineRequest(
                name=self.resource_name,
                input=kwargs,
                class_method=method_name,
            ),
        )
        for chunk in response:
            for parsed_json in _utils.yield_parsed_json(chunk):
                if parsed_json is not None:
                    yield parsed_json

    return _method


def _wrap_async_stream_query_operation(
    *, method_name: str
) -> Callable[..., AsyncIterable[Any]]:
    """Wraps an Agent Engine method, creating an async callable for `stream_query` API.

    This function creates a callable object that executes the specified
    Agent Engine method using the `stream_query` API.  It handles the
    creation of the API request and the processing of the API response.

    Args:
        method_name: The name of the Agent Engine method to call.
        doc: Documentation string for the method.

    Returns:
        A callable object that executes the method on the Agent Engine via
        the `stream_query` API.
    """

    async def _method(self, **kwargs) -> AsyncIterable[Any]:
        response = self.execution_api_client.stream_query_reasoning_engine(
            request=aip_types.StreamQueryReasoningEngineRequest(
                name=self.resource_name,
                input=kwargs,
                class_method=method_name,
            ),
        )
        for chunk in response:
            for parsed_json in _utils.yield_parsed_json(chunk):
                if parsed_json is not None:
                    yield parsed_json

    return _method


def _wrap_bidi_stream_query_operation(
    *, method_name: str
) -> Callable[..., AsyncIterable[Any]]:
    """Wraps an Agent Engine method, creating an async callable for `bidi_stream_query` API.

    This function creates a callable object that executes the specified
    Agent Engine method using the `bidi_stream_query` API.  It handles the
    creation of the API request and the processing of the API response.

    Args:
        method_name: The name of the Agent Engine method to call.

    Returns:
        A callable object that executes the method on the Agent Engine via
        the `bidi_stream_query` API.
    """

    async def _method(self, **kwargs) -> AsyncIterable[Any]:
        # Agent Engine bidi streaming query execution should use GenAI SDK Agent
        # Engine live API client directly.
        raise NotImplementedError(
            f"{method_name} is not implemented, please use GenAI SDK Agent "
            "Enginve live API client instead."
        )


def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list]:
    """Wraps an Agent Engine method, creating a callable for A2A API.

    Args:
        method_name: The name of the Agent Engine method to call.
        agent_card: The agent card to use for the A2A API call.
            Example:
                {'additionalInterfaces': None,
                'capabilities': {'extensions': None,
                'pushNotifications': None,
                'stateTransitionHistory': None,
                'streaming': False},
                'defaultInputModes': ['text'],
                'defaultOutputModes': ['text'],
                'description': (
                    'A helpful assistant agent that can answer questions.'
                ),
                'documentationUrl': None,
                'iconUrl': None,
                'name': 'Q&A Agent',
                'preferredTransport': 'JSONRPC',
                'protocolVersion': '0.3.0',
                'provider': None,
                'security': None,
                'securitySchemes': None,
                'signatures': None,
                'skills': [{
                    'description': (
                        'A helpful assistant agent that can answer questions.'
                    ),
                    'examples': ['Who is leading 2025 F1 Standings?',
                        'Where can i find an active volcano?'],
                    'id': 'question_answer',
                'inputModes': None,
                'name': 'Q&A Agent',
                'outputModes': None,
                'security': None,
                'tags': ['Question-Answer']}],
                'supportsAuthenticatedExtendedCard': True,
                'url': 'http://localhost:8080/',
                'version': '1.0.0'}
    Returns:
        A callable object that executes the method on the Agent Engine via
        the A2A API.
    """

    async def _method(self, **kwargs) -> Any:
        """Wraps an Agent Engine method, creating a callable for A2A API."""
        a2a_agent_card = AgentCard(**json.loads(agent_card))

        # A2A + AE integration currently only supports Rest API.
        if (
            a2a_agent_card.preferred_transport
            and a2a_agent_card.preferred_transport != TransportProtocol.http_json
        ):
            raise ValueError(
                "Only HTTP+JSON is supported for preferred transport on agent card "
            )

        # Set preferred transport to HTTP+JSON if not set.
        if not hasattr(a2a_agent_card, "preferred_transport"):
            a2a_agent_card.preferred_transport = TransportProtocol.http_json

        # AE cannot support streaming yet. Turn off streaming for now.
        if a2a_agent_card.capabilities and a2a_agent_card.capabilities.streaming:
            raise ValueError(
                "Streaming is not supported in Agent Engine, please change "
                "a2a_agent_card.capabilities.streaming to False."
            )

        if not hasattr(a2a_agent_card.capabilities, "streaming"):
            a2a_agent_card.capabilities.streaming = False

        # agent_card is set on the class_methods before set_up is invoked.
        # Ensure that the agent_card url is set correctly before the client is created.
        a2a_agent_card.url = f"https://{initializer.global_config.api_endpoint}/v1beta1/{self.resource_name}/a2a"

        # Using a2a client, inject the auth token from the global config.
        config = ClientConfig(
            supported_transports=[
                TransportProtocol.http_json,
            ],
            use_client_preference=True,
            httpx_client=httpx.AsyncClient(
                headers={
                    "Authorization": (
                        f"Bearer {initializer.global_config.credentials.token}"
                    )
                }
            ),
        )
        factory = ClientFactory(config)
        client = factory.create(a2a_agent_card)

        # kokoro job uses python 3.9, replaced match with if else.
        if method_name == "on_message_send":
            response = client.send_message(Message(**kwargs))
            chunks = []
            async for chunk in response:
                chunks.append(chunk)
            return chunks
        elif method_name == "on_get_task":
            response = await client.get_task(TaskQueryParams(**kwargs))
        elif method_name == "on_cancel_task":
            response = await client.cancel_task(TaskIdParams(**kwargs))
        elif method_name == "handle_authenticated_agent_card":
            response = await client.get_card()
        else:
            raise ValueError(f"Unknown method name: {method_name}")

        return response

    return _method


def _unregister_api_methods(
    obj: "AgentEngine", operation_schemas: Sequence[_utils.JsonDict]
):
    """Unregisters Agent Engine API methods based on operation schemas.

    This function iterates through operation schemas provided by the
    AgentEngine object.  Each schema defines an API mode and method name.
    It dynamically unregisters methods on the AgentEngine object. This
    should only be used when updating the object.

    Args:
        obj: The AgentEngine object to augment with API methods.
        operation_schemas: The operation schemas to use for method unregistration.
    """
    for operation_schema in operation_schemas:
        if "name" in operation_schema:
            method_name = operation_schema.get("name")
            if hasattr(obj, method_name):
                delattr(obj, method_name)


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
    AgentEngine object.  Each schema defines an API mode and method name.
    It dynamically creates and registers methods on the AgentEngine object
    to handle API calls based on the specified API mode.
    Currently, only standard API mode `` is supported.

    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
        is missing any required fields (e.g. `api_mode` or `name`).
    """
    for operation_schema in obj.operation_schemas():
        if _MODE_KEY_IN_SCHEMA not in operation_schema:
            raise ValueError(
                f"Operation schema {operation_schema} does not"
                f" contain an `{_MODE_KEY_IN_SCHEMA}` field."
            )
        api_mode = operation_schema.get(_MODE_KEY_IN_SCHEMA)
        if _METHOD_NAME_KEY_IN_SCHEMA not in operation_schema:
            raise ValueError(
                f"Operation schema {operation_schema} does not"
                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",
            _DEFAULT_METHOD_DOCSTRING_TEMPLATE.format(
                method_name=method_name,
                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,
            _BIDI_STREAM_API_MODE: _wrap_bidi_stream_query_operation,
            _A2A_EXTENSION_MODE: _wrap_a2a_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:
            supported_api_modes = ", ".join(
                f"`{mode}`" for mode in sorted(_wrap_operation_map.keys())
            )
            raise ValueError(
                f"Unsupported api mode: `{api_mode}`,"
                f" Supported modes are: {supported_api_modes}."
            )

        # Bind the method to the object.
        if api_mode == _A2A_EXTENSION_MODE:
            agent_card = operation_schema.get(_A2A_AGENT_CARD)
            method = _wrap_operation(method_name=method_name, agent_card=agent_card)
        else:
            method = _wrap_operation(method_name=method_name)
        method.__name__ = method_name
        method.__doc__ = method_description
        setattr(obj, method_name, types.MethodType(method, obj))


def _get_registered_operations(
    agent_engine: _AgentEngineInterface,
) -> Dict[str, List[str]]:
    """Retrieves registered operations for a AgentEngine."""
    if isinstance(agent_engine, OperationRegistrable):
        return agent_engine.register_operations()

    operations = {}
    if isinstance(agent_engine, Queryable):
        operations[_STANDARD_API_MODE] = [_DEFAULT_METHOD_NAME]
    if isinstance(agent_engine, AsyncQueryable):
        operations[_ASYNC_API_MODE] = [_DEFAULT_ASYNC_METHOD_NAME]
    if isinstance(agent_engine, StreamQueryable):
        operations[_STREAM_API_MODE] = [_DEFAULT_STREAM_METHOD_NAME]
    if isinstance(agent_engine, AsyncStreamQueryable):
        operations[_ASYNC_STREAM_API_MODE] = [_DEFAULT_ASYNC_STREAM_METHOD_NAME]
    return operations


def _generate_class_methods_spec_or_raise(
    *,
    agent_engine: _AgentEngineInterface,
    operations: Dict[str, List[str]],
    logger: base.Logger = _LOGGER,
) -> List[proto.Message]:
    """Generates a ReasoningEngineSpec based on the registered operations.

    Args:
        agent_engine: The AgentEngine instance.
        operations: A dictionary of API modes and method names.

    Returns:
        A list of ReasoningEngineSpec.ClassMethod messages.

    Raises:
        ValueError: If a method defined in `register_operations` is not found on
        the AgentEngine.
    """
    if isinstance(agent_engine, ModuleAgent):
        # We do a dry-run of setting up the agent engine to have the operations
        # needed for registration.
        agent_engine = agent_engine.clone()
        try:
            agent_engine.set_up()
        except Exception as e:
            raise ValueError(
                f"Failed to set up agent engine {agent_engine}: {e}"
            ) from e
    class_methods_spec = []
    for mode, method_names in operations.items():
        for method_name in method_names:
            if mode == _BIDI_STREAM_API_MODE:
                _LOGGER.warning(
                    "Bidi stream API mode is not supported yet in Vertex SDK, "
                    "please use the GenAI SDK instead. Skipping "
                    f"method {method_name}."
                )
                continue
            if not hasattr(agent_engine, method_name):
                raise ValueError(
                    f"Method `{method_name}` defined in `register_operations`"
                    " not found on AgentEngine."
                )

            method = getattr(agent_engine, method_name)
            try:
                schema_dict = _utils.generate_schema(method, schema_name=method_name)
            except Exception as e:
                logger.warning(f"failed to generate schema for {method_name}: {e}")
                continue

            class_method = _utils.to_proto(schema_dict)
            class_method[_MODE_KEY_IN_SCHEMA] = mode
            # A2A agent card is a special case, when running in A2A mode,
            if hasattr(agent_engine, "agent_card"):
                class_method[_A2A_AGENT_CARD] = getattr(
                    agent_engine, "agent_card"
                ).model_dump_json()
            class_methods_spec.append(class_method)

    return class_methods_spec


def _class_methods_to_class_methods_spec(
    class_methods: List[dict[str, Any]],
) -> List[proto.Message]:
    """Converts a list of class methods to a list of ReasoningEngineSpec.ClassMethod messages."""
    return [_utils.to_proto(class_method) for class_method in class_methods]