_agent_engines_utils.py

# 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.
#
"""Utility functions for agent engines."""

import abc
import asyncio
import base64
import dataclasses
from importlib import metadata as importlib_metadata
import inspect
import io
import json
import logging
import os
import re
import sys
import tarfile
import time
import types
import typing
from typing import (
    Any,
    AsyncIterator,
    Callable,
    Coroutine,
    Dict,
    Iterator,
    List,
    Mapping,
    Optional,
    Protocol,
    Sequence,
    Set,
    TypedDict,
    Union,
)

import httpx

import proto

from google.api_core import exceptions
from google.genai import types as google_genai_types
from google.protobuf import struct_pb2
from google.protobuf import json_format

from . import types as genai_types


if sys.version_info >= (3, 10):
    from typing import TypeAlias
else:
    from typing_extensions import TypeAlias


try:
    _BUILTIN_MODULE_NAMES: Sequence[str] = sys.builtin_module_names
except AttributeError:
    _BUILTIN_MODULE_NAMES: Sequence[str] = []  # type: ignore[no-redef]

try:
    _PACKAGE_DISTRIBUTIONS: Mapping[str, Sequence[str]] = (
        importlib_metadata.packages_distributions()
    )
except AttributeError:
    _PACKAGE_DISTRIBUTIONS: Mapping[str, Sequence[str]] = {}  # type: ignore[no-redef]

try:
    # sys.stdlib_module_names is available from Python 3.10 onwards.
    _STDLIB_MODULE_NAMES: frozenset[str] = sys.stdlib_module_names
except AttributeError:
    _STDLIB_MODULE_NAMES: frozenset[str] = frozenset()  # type: ignore[no-redef]


if typing.TYPE_CHECKING:
    from google.cloud import storage  # type: ignore[attr-defined]

    _StorageBucket: TypeAlias = storage.Bucket
else:
    try:
        from google.cloud import storage  # type: ignore[attr-defined]

        _StorageBucket: type[Any] = storage.Bucket
    except (ImportError, AttributeError):
        _StorageBucket: type[Any] = Any  # type: ignore[no-redef]


if typing.TYPE_CHECKING:
    import packaging

    _SpecifierSet = packaging.specifiers.SpecifierSet
else:
    try:
        import packaging

        _SpecifierSet: type[Any] = packaging.specifiers.SpecifierSet
    except (ImportError, AttributeError):
        _SpecifierSet: type[Any] = Any  # type: ignore[no-redef]


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
try:
    from autogen.agentchat import chat

    AutogenChatResult = chat.ChatResult
except ImportError:
    AutogenChatResult = Any
try:
    from autogen.io import run_response

    AutogenRunResponse = run_response.RunResponse
except ImportError:
    AutogenRunResponse = Any
try:
    from llama_index.core.base.response import schema as llama_index_schema
    from llama_index.core.base.llms import types as llama_index_types

    LlamaIndexResponse = llama_index_schema.Response
    LlamaIndexBaseModel = llama_index_schema.BaseModel
    LlamaIndexChatResponse = llama_index_types.ChatResponse
except ImportError:
    LlamaIndexResponse = Any
    LlamaIndexBaseModel = Any
    LlamaIndexChatResponse = Any
try:
    import pydantic

    BaseModel = pydantic.BaseModel
except ImportError:
    BaseModel = Any

JsonDict = Dict[str, Any]

_ACTIONS_KEY = "actions"
_ACTION_APPEND = "append"
_AGENT_FRAMEWORK_ATTR = "agent_framework"
_ASYNC_API_MODE = "async"
_ASYNC_STREAM_API_MODE = "async_stream"
_BIDI_STREAM_API_MODE = "bidi_stream"
_BASE_MODULES = set(_BUILTIN_MODULE_NAMES).union(_STDLIB_MODULE_NAMES)
_BLOB_FILENAME = "agent_engine.pkl"
_DEFAULT_AGENT_FRAMEWORK = "custom"
_SUPPORTED_AGENT_FRAMEWORKS = frozenset(
    [
        "google-adk",
        "langchain",
        "langgraph",
        "ag2",
        "llama-index",
        "custom",
        "a2a",
    ]
)
_DEFAULT_ASYNC_METHOD_NAME = "async_query"
_DEFAULT_ASYNC_METHOD_RETURN_TYPE = "Coroutine[Any]"
_DEFAULT_ASYNC_STREAM_METHOD_NAME = "async_stream_query"
_DEFAULT_ASYNC_STREAM_METHOD_RETURN_TYPE = "AsyncIterable[Any]"
_DEFAULT_GCS_DIR_NAME = "agent_engine"
_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.
"""
_DEFAULT_METHOD_NAME = "query"
_DEFAULT_METHOD_RETURN_TYPE = "dict[str, Any]"
_DEFAULT_STREAM_METHOD_RETURN_TYPE = "Iterable[Any]"
_DEFAULT_REQUIRED_PACKAGES = frozenset(["cloudpickle", "pydantic"])
_DEFAULT_STREAM_METHOD_NAME = "stream_query"
_DEFAULT_BIDI_STREAM_METHOD_NAME = "bidi_stream_query"
_EXTRA_PACKAGES_FILE = "dependencies.tar.gz"
_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}"
)
_INSTALLATION_SUBDIR = "installation_scripts"
_METHOD_NAME_KEY_IN_SCHEMA = "name"
_MODE_KEY_IN_SCHEMA = "api_mode"
_REQUIREMENTS_FILE = "requirements.txt"
_STANDARD_API_MODE = ""
_STREAM_API_MODE = "stream"
_A2A_EXTENSION_MODE = "a2a_extension"
_A2A_AGENT_CARD = "a2a_agent_card"
_WARNINGS_KEY = "warnings"
_WARNING_MISSING = "missing"
_WARNING_INCOMPATIBLE = "incompatible"

_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,
}


logger = logging.getLogger("agentplatform_genai.agentengines")


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

    @abc.abstractmethod
    def query(self, **kwargs):  # type: ignore[no-untyped-def]
        """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):  # type: ignore[no-untyped-def]
        """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) -> AsyncIterator[Any]:  # type: ignore[no-untyped-def]
        """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) -> Iterator[Any]:  # type: ignore[no-untyped-def]
        """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, input_queue: asyncio.Queue[Any]
    ) -> AsyncIterator[Any]:
        """Stream requests and responses to serve the user queries."""


@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: Any) -> dict[str, list[str]]:
        """Register the user provided operations (modes and methods)."""
        pass


if typing.TYPE_CHECKING:
    from google.adk.agents import BaseAgent

    ADKAgent: TypeAlias = BaseAgent
else:
    try:
        from google.adk.agents import BaseAgent

        ADKAgent: Optional[TypeAlias] = BaseAgent
    except (ImportError, AttributeError):
        ADKAgent = None  # type: ignore[no-redef]

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


class _ModuleAgentAttributes(TypedDict, total=False):
    module_name: str
    agent_name: str
    register_operations: Dict[str, list[str]]
    sys_paths: Optional[Sequence[str]]
    agent: _AgentEngineInterface


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, list[str]],
        sys_paths: Optional[Sequence[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, list[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._tmpl_attrs: _ModuleAgentAttributes = {
            "module_name": module_name,
            "agent_name": agent_name,
            "register_operations": register_operations,
            "sys_paths": sys_paths,
        }

    def clone(self) -> "ModuleAgent":
        """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"),
        )

    def register_operations(self, **kwargs: Any) -> dict[str, list[str]]:
        reg_operations = self._tmpl_attrs.get("register_operations")
        if reg_operations is None:
            raise ValueError("Register operations is not set.")
        return reg_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.
        """
        sys_paths = self._tmpl_attrs.get("sys_paths")
        if isinstance(sys_paths, Sequence):
            import sys

            for sys_path in 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
        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=agent, operation=operation)
                setattr(self, operation, types.MethodType(op, self))


class _RequirementsValidationActions(TypedDict):
    append: Set[str]


class _RequirementsValidationWarnings(TypedDict):
    missing: Set[str]
    incompatible: Set[str]


class _RequirementsValidationResult(TypedDict):
    warnings: _RequirementsValidationWarnings
    actions: _RequirementsValidationActions


AgentEngineOperationUnion = Union[
    genai_types.AgentEngineOperation,
    genai_types.AgentEngineMemoryOperation,
    genai_types.AgentEngineGenerateMemoriesOperation,
]


class GetOperationFunction(Protocol):
    def __call__(
        self, *, operation_name: str, **kwargs: Any
    ) -> AgentEngineOperationUnion:
        pass


class GetAsyncOperationFunction(Protocol):
    async def __call__(
        self, *, operation_name: str, **kwargs: Any
    ) -> AgentEngineOperationUnion:
        pass


def _get_reasoning_engine_id(operation_name: str = "", resource_name: str = "") -> str:
    """Returns reasoning engine ID from operation name or resource name."""
    if not resource_name and not operation_name:
        raise ValueError("Resource name or operation name cannot be empty.")

    if resource_name:
        match = re.match(
            r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)$",
            resource_name,
        )
        if match:
            return match.group(1)
        else:
            raise ValueError(
                "Failed to parse reasoning engine ID from resource name: "
                f"`{resource_name}`"
            )

    if not operation_name:
        raise ValueError("Operation name cannot be empty.")

    match = re.match(
        r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)/operations/[^/]+$",
        operation_name,
    )
    if match:
        return match.group(1)
    raise ValueError(
        "Failed to parse reasoning engine ID from operation name: "
        f"`{operation_name}`"
    )


async def _await_async_operation(
    *,
    operation_name: str,
    get_operation_fn: GetAsyncOperationFunction,
    poll_interval_seconds: float = 10,
) -> Any:
    """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 (float):
            The number of seconds to wait between each poll.
        get_operation_fn (Callable[[str], Awaitable[Any]]):
            Optional. The async function to use for getting the operation. If not
            provided, `self._get_agent_operation` will be used.

    Returns:
        The operation that has completed (i.e. `operation.done==True`).
    """
    operation = await get_operation_fn(operation_name=operation_name)
    while not operation.done:
        await asyncio.sleep(poll_interval_seconds)
        operation = await get_operation_fn(operation_name=operation.name)

    return operation


def _await_operation(
    *,
    operation_name: str,
    get_operation_fn: GetOperationFunction,
    poll_interval_seconds: float = 10,
) -> Any:
    """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 (float):
            The number of seconds to wait between each poll.
        get_operation_fn (Callable[[str], Any]):
            Optional. The function to use for getting the operation. If not
            provided, `self._get_agent_operation` will be used.

    Returns:
        The operation that has completed (i.e. `operation.done==True`).
    """
    operation = get_operation_fn(operation_name=operation_name)
    while not operation.done:
        time.sleep(poll_interval_seconds)
        operation = get_operation_fn(operation_name=operation.name)

    return operation


def _compare_requirements(
    *,
    requirements: Mapping[str, str],
    constraints: Union[Sequence[str], Mapping[str, Optional["_SpecifierSet"]]],
    required_packages: Optional[Iterator[str]] = None,
) -> _RequirementsValidationResult:
    """Compares the requirements with the constraints.

    Args:
        requirements (Mapping[str, str]):
            Required. The packages (and their versions) to compare with the constraints.
            This is assumed to be the result of `scan_requirements`.
        constraints (Union[Sequence[str], Mapping[str, SpecifierSet]]):
            Required. The package constraints to compare against. This is assumed
            to be the result of `parse_constraints`.
        required_packages (Iterator[str]):
            Optional. The set of packages that are required to be in the
            constraints. It defaults to the set of packages that are required
            for deployment on Agent Engine.

    Returns:
        dict[str, dict[str, Any]]: The comparison result as a dictionary containing:
            * warnings:
                * missing: The set of packages that are not in the constraints.
                * incompatible: The set of packages that are in the constraints
                    but have versions that are not in the constraint specifier.
            * actions:
                * append: The set of packages that are not in the constraints
                    but should be appended to the constraints.
    """
    packaging_version = _import_packaging_version_or_raise()
    if required_packages is None:
        required_packages = _DEFAULT_REQUIRED_PACKAGES  # type: ignore[assignment]
    result = _RequirementsValidationResult(
        warnings=_RequirementsValidationWarnings(missing=set(), incompatible=set()),
        actions=_RequirementsValidationActions(append=set()),
    )
    if isinstance(constraints, list):
        constraints = _parse_constraints(constraints=constraints)
    for package, package_version in requirements.items():
        if package not in constraints:
            result[_WARNINGS_KEY][_WARNING_MISSING].add(package)  # type: ignore[literal-required]
            if package in required_packages:  # type: ignore[operator]
                result[_ACTIONS_KEY][_ACTION_APPEND].add(  # type: ignore[literal-required]
                    f"{package}=={package_version}"
                )
            continue
        if package_version:
            package_specifier = constraints[package]  # type: ignore[call-overload]
            if not package_specifier:
                continue
            if packaging_version.Version(package_version) not in package_specifier:
                result[_WARNINGS_KEY][_WARNING_INCOMPATIBLE].add(  # type: ignore[literal-required]
                    f"{package}=={package_version} (required: {str(package_specifier)})"
                )
    return result


def _generate_class_methods_spec_or_raise(
    *,
    agent: _AgentEngineInterface,
    operations: Dict[str, List[str]],
) -> List[proto.Message]:
    """Generates a ReasoningEngineSpec based on the registered operations.

    Args:
        agent: 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, ModuleAgent):
        # We do a dry-run of setting up the agent engine to have the operations
        # needed for registration.
        agent: ModuleAgent = agent.clone()  # type: ignore[no-redef]
        try:
            agent.set_up()
        except Exception as e:
            raise ValueError(f"Failed to set up agent {agent}: {e}") from e
    class_methods_spec = []
    for mode, method_names in operations.items():
        for method_name in method_names:
            if not hasattr(agent, method_name):
                raise ValueError(
                    f"Method `{method_name}` defined in `register_operations`"
                    " not found on agent."
                )

            method = getattr(agent, method_name)
            try:
                schema_dict = _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 = _to_proto(schema_dict)
            class_method[_MODE_KEY_IN_SCHEMA] = mode
            if hasattr(agent, "agent_card"):
                card = getattr(agent, "agent_card")
                if hasattr(card, "model_dump_json"):
                    class_method[_A2A_AGENT_CARD] = card.model_dump_json()
                elif hasattr(card, "DESCRIPTOR"):
                    class_method[_A2A_AGENT_CARD] = json_format.MessageToJson(card)
                elif isinstance(card, str):
                    class_method[_A2A_AGENT_CARD] = card
                else:
                    class_method[_A2A_AGENT_CARD] = json.dumps(card)
            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 [_to_proto(class_method) for class_method in class_methods]


def _is_pydantic_serializable(param: inspect.Parameter) -> bool:
    """Checks if the parameter is pydantic serializable."""

    if param.annotation == inspect.Parameter.empty:
        return True

    if "ForwardRef" in repr(param.annotation):
        return True

    if isinstance(param.annotation, str):
        return False

    pydantic = _import_pydantic_or_raise()
    try:
        pydantic.TypeAdapter(param.annotation)
        return True
    except Exception:
        return False


def _generate_schema(
    f: Callable[..., Any],
    *,
    schema_name: Optional[str] = None,
    descriptions: Mapping[str, str] = {},
    required: Sequence[str] = [],
) -> Dict[str, Any]:
    """Generates the OpenAPI Schema for a callable object.

    Only positional and keyword arguments of the function `f` will be supported
    in the OpenAPI Schema that is generated. I.e. `*args` and `**kwargs` will
    not be present in the OpenAPI schema returned from this function. For those
    cases, you can either include it in the docstring for `f`, or modify the
    OpenAPI schema returned from this function to include additional arguments.

    Args:
        f (Callable):
            Required. The function to generate an OpenAPI Schema for.
        schema_name (str):
            Optional. The name for the OpenAPI schema. If unspecified, the name
            of the Callable will be used.
        descriptions (Mapping[str, str]):
            Optional. A `{name: description}` mapping for annotating input
            arguments of the function with user-provided descriptions. It
            defaults to an empty dictionary (i.e. there will not be any
            description for any of the inputs).
        required (Sequence[str]):
            Optional. For the user to specify the set of required arguments in
            function calls to `f`. If specified, it will be automatically
            inferred from `f`.

    Returns:
        dict[str, Any]: The OpenAPI Schema for the function `f` in JSON format.
    """
    pydantic = _import_pydantic_or_raise()
    defaults = dict(inspect.signature(f).parameters)
    fields_dict = {
        name: (
            # 1. We infer the argument type here: use Any rather than None so
            # it will not try to auto-infer the type based on the default value.
            (
                param.annotation
                if param.annotation != inspect.Parameter.empty
                and "ForwardRef" not in repr(param.annotation)
                else Any
            ),
            pydantic.Field(
                # 2. We do not support default values for now.
                # default=(
                #     param.default if param.default != inspect.Parameter.empty
                #     else None
                # ),
                # 3. We support user-provided descriptions.
                description=descriptions.get(name, None),
            ),
        )
        for name, param in defaults.items()
        # We do not support *args or **kwargs
        if param.kind
        in (
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
            inspect.Parameter.KEYWORD_ONLY,
            inspect.Parameter.POSITIONAL_ONLY,
        )
        # For a bidi endpoint, it requires an asyncio.Queue as the input, but
        # it is not JSON serializable. We hence exclude it from the schema.
        and param.annotation != asyncio.Queue and _is_pydantic_serializable(param)
    }
    parameters = pydantic.create_model(f.__name__, **fields_dict).schema()
    # Postprocessing
    # 4. Suppress unnecessary title generation:
    #    * https://github.com/pydantic/pydantic/issues/1051
    #    * http://cl/586221780
    parameters.pop("title", "")
    for name, function_arg in parameters.get("properties", {}).items():
        function_arg.pop("title", "")
        annotation = defaults[name].annotation
        # 5. Nullable fields:
        #     * https://github.com/pydantic/pydantic/issues/1270
        #     * https://stackoverflow.com/a/58841311
        #     * https://github.com/pydantic/pydantic/discussions/4872
        if typing.get_origin(annotation) is Union and type(None) in typing.get_args(
            annotation
        ):
            # for "typing.Optional" arguments, function_arg might be a
            # dictionary like
            #
            #   {'anyOf': [{'type': 'integer'}, {'type': 'null'}]
            for schema in function_arg.pop("anyOf", []):
                schema_type = schema.get("type")
                if schema_type and schema_type != "null":
                    function_arg["type"] = schema_type
                    break
            function_arg["nullable"] = True
    # 6. Annotate required fields.
    if required:
        # We use the user-provided "required" fields if specified.
        parameters["required"] = required
    else:
        # Otherwise we infer it from the function signature.
        parameters["required"] = [
            k
            for k in defaults
            if (
                defaults[k].default == inspect.Parameter.empty
                and defaults[k].kind
                in (
                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
                    inspect.Parameter.KEYWORD_ONLY,
                    inspect.Parameter.POSITIONAL_ONLY,
                )
            )
        ]
    schema = dict(name=f.__name__, description=f.__doc__, parameters=parameters)
    if schema_name:
        schema["name"] = schema_name
    return schema


def _get_agent_framework(
    *,
    agent_framework: Optional[str],
    agent: _AgentEngineInterface,
) -> Union[str, Any]:
    """Gets the agent framework to use.

    The agent framework is determined in the following order of priority:
    1. The `agent_framework` passed to this function.
    2. The `agent_framework` attribute on the `agent` object.
    3. The default framework, "custom".

    Args:
        agent_framework (str):
            The agent framework provided by the user.
        agent (_AgentEngineInterface):
            The agent engine instance.

    Returns:
        str: The name of the agent framework to use.
    """
    if agent_framework is not None and agent_framework in _SUPPORTED_AGENT_FRAMEWORKS:
        logger.info(f"Using agent framework: {agent_framework}")
        return agent_framework
    if hasattr(agent, _AGENT_FRAMEWORK_ATTR):
        agent_framework_attr = getattr(agent, _AGENT_FRAMEWORK_ATTR)
        if (
            agent_framework_attr is not None
            and isinstance(agent_framework_attr, str)
            and agent_framework_attr in _SUPPORTED_AGENT_FRAMEWORKS
        ):
            logger.info(f"Using agent framework: {agent_framework_attr}")
            return agent_framework_attr
    logger.info(
        f"The provided agent framework {agent_framework} is not supported."
        f" Defaulting to {_DEFAULT_AGENT_FRAMEWORK}."
    )
    return _DEFAULT_AGENT_FRAMEWORK


def _get_gcs_bucket(
    *,
    project: str,
    location: str,
    staging_bucket: str,
    credentials: Optional[Any] = None,
) -> _StorageBucket:
    """Gets or creates the GCS bucket."""
    storage = _import_cloud_storage_or_raise()
    storage_client = storage.Client(project=project, credentials=credentials)
    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 _get_registered_operations(
    *,
    agent: _AgentEngineInterface,
) -> dict[str, list[str]]:
    """Retrieves registered operations for a AgentEngine."""
    if isinstance(agent, OperationRegistrable):
        return agent.register_operations()

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


def _import_cloudpickle_or_raise() -> types.ModuleType:
    """Tries to import the cloudpickle module."""
    try:
        import cloudpickle  # noqa:F401
    except ImportError as e:
        raise ImportError(
            "cloudpickle is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        ) from e
    return cloudpickle  # type: ignore[no-any-return]


def _import_cloud_storage_or_raise() -> types.ModuleType:
    """Tries to import the Cloud Storage module."""
    try:
        from google.cloud import storage  # type: ignore[attr-defined]
    except ImportError as e:
        raise ImportError(
            "Cloud Storage is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        ) from e
    return storage  # type: ignore[no-any-return]


def _import_packaging_requirements_or_raise() -> types.ModuleType:
    """Tries to import the packaging.requirements module."""
    try:
        from packaging import requirements
    except ImportError as e:
        raise ImportError(
            "packaging.requirements is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        ) from e
    return requirements


def _import_packaging_version_or_raise() -> types.ModuleType:
    """Tries to import the packaging.requirements module."""
    try:
        from packaging import version
    except ImportError as e:
        raise ImportError(
            "packaging.version is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        ) from e
    return version


def _import_pydantic_or_raise() -> types.ModuleType:
    """Tries to import the pydantic module."""
    try:
        import pydantic

        _ = pydantic.Field
    except AttributeError:
        from pydantic import v1 as pydantic  # type: ignore[no-redef]
    except ImportError as e:
        raise ImportError(
            "pydantic is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        ) from e
    return pydantic


def _parse_constraints(
    *,
    constraints: Sequence[str],
) -> Mapping[str, Optional["_SpecifierSet"]]:
    """Parses a list of constraints into a dict of requirements.

    Args:
        constraints (list[str]):
            Required. The list of package requirements to parse. This is assumed
            to come from the `requirements.txt` file.

    Returns:
        dict[str, SpecifierSet]: The specifiers for each package.
    """
    requirements = _import_packaging_requirements_or_raise()
    result: Dict[str, Optional[_SpecifierSet]] = {}
    for constraint in constraints:
        try:
            if constraint.endswith(".whl"):
                constraint = os.path.basename(constraint)
            requirement = requirements.Requirement(constraint)
        except Exception as e:
            logger.warning(f"Failed to parse constraint: {constraint}. Exception: {e}")
            continue
        result[requirement.name] = requirement.specifier or None
    return result


def _prepare(
    *,
    agent: Optional[_AgentEngineInterface],
    requirements: Optional[Sequence[str]],
    extra_packages: Optional[Sequence[str]],
    project: str,
    location: str,
    staging_bucket: str,
    gcs_dir_name: str,
    credentials: Optional[Any] = None,
) -> 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: 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.
        credentials: The credentials to use for the storage client.
    """
    if agent is None:
        return
    gcs_bucket = _get_gcs_bucket(
        project=project,
        location=location,
        staging_bucket=staging_bucket,
        credentials=credentials,
    )
    _upload_agent_engine(
        agent=agent,
        gcs_bucket=gcs_bucket,
        gcs_dir_name=gcs_dir_name,
    )
    if requirements is not None:
        _upload_requirements(
            requirements=requirements,
            gcs_bucket=gcs_bucket,
            gcs_dir_name=gcs_dir_name,
        )
    if extra_packages is not None:
        _upload_extra_packages(
            extra_packages=extra_packages,
            gcs_bucket=gcs_bucket,
            gcs_dir_name=gcs_dir_name,
        )


def _register_api_methods_or_raise(
    *,
    agent_engine: genai_types.AgentEngine | genai_types.AgentEngineRuntimeRevision,
    wrap_operation_fn: Optional[
        dict[str, Callable[[str, str], Callable[..., Any]]]
    ] = None,
) -> None:
    """Registers Agent Engine API methods based on operation schemas.

    This function iterates through operation schemas provided by the
    `agent_engine`.  Each schema defines an API mode and method name.
    It dynamically creates and registers methods on the `agent_engine`
    to handle API calls based on the specified API mode.
    Currently, only standard API mode `` is supported.

    Args:
        agent_engine: The AgentEngine 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`).
    """
    operation_schemas = agent_engine.operation_schemas()
    if not operation_schemas:
        return
    for operation_schema in 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)
        # For bidi stream api mode, we don't need to wrap the operation.
        if api_mode == _BIDI_STREAM_API_MODE:
            continue

        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)
        if not isinstance(method_name, str):
            raise ValueError(
                "Operation schema has a non-string value for"
                f" `{_METHOD_NAME_KEY_IN_SCHEMA}`: {method_name}"
            )
        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,
            _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]  # type: ignore[assignment]
        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
            )  # type: ignore[call-arg]
        else:
            method = _wrap_operation(method_name=method_name)  # type: ignore[call-arg]
        method.__name__ = method_name
        if method_description and isinstance(method_description, str):
            method.__doc__ = method_description
        setattr(agent_engine, method_name, types.MethodType(method, agent_engine))


def _scan_requirements(
    *,
    obj: Any,
    ignore_modules: Optional[Sequence[str]] = None,
    package_distributions: Optional[Mapping[str, Sequence[str]]] = None,
    inspect_getmembers_kwargs: Optional[Mapping[str, Any]] = None,
) -> Mapping[str, str]:
    """Scans the object for modules and returns the requirements discovered.

    This is not a comprehensive scan of the object, and only detects for common
    cases based on the members of the object returned by `dir(obj)`.

    Args:
        obj (Any):
            Required. The object to scan for package requirements.
        ignore_modules (Sequence[str]):
            Optional. The set of modules to ignore. It defaults to the set of
            built-in and stdlib modules.
        package_distributions (Mapping[str, Sequence[str]]):
            Optional. The mapping of module names to the set of packages that
            contain them. It defaults to the set of packages from
            `importlib_metadata.packages_distributions()`.
        inspect_getmembers_kwargs (Mapping[str, Any]):
            Optional. The keyword arguments to pass to `inspect.getmembers`. It
            defaults to an empty dictionary.

    Returns:
        Sequence[str]: The list of requirements that were discovered.
    """
    if ignore_modules is None:
        ignore_modules = _BASE_MODULES  # type: ignore[assignment]
    if package_distributions is None:
        package_distributions = _PACKAGE_DISTRIBUTIONS
    modules_found = set(_DEFAULT_REQUIRED_PACKAGES)
    inspect_getmembers_kwargs = inspect_getmembers_kwargs or {}
    for _, attr in inspect.getmembers(obj, **inspect_getmembers_kwargs):
        if not attr or inspect.isbuiltin(attr) or not hasattr(attr, "__module__"):
            continue
        module_name = (attr.__module__ or "").split(".")[0]
        if module_name and module_name not in ignore_modules:  # type: ignore[operator]
            for module in package_distributions.get(module_name, []):
                modules_found.add(module)
    return {module: importlib_metadata.version(module) for module in modules_found}


def _to_dict(message: proto.Message) -> Dict[str, Any]:
    """Converts the contents of the protobuf message to JSON format.

    Args:
        message (proto.Message):
            Required. The proto message to be converted to a JSON dictionary.

    Returns:
        dict[str, Any]: A dictionary containing the contents of the proto.
    """
    try:
        # Best effort attempt to convert the message into a JSON dictionary.
        result: Dict[str, Any] = json.loads(
            json_format.MessageToJson(
                message._pb,
                preserving_proto_field_name=True,
            )
        )
    except AttributeError:
        result: Dict[str, Any] = json.loads(  # type: ignore[no-redef]
            json_format.MessageToJson(
                message,
                preserving_proto_field_name=True,
            )
        )
    return result


def _to_proto(
    obj: Union[Dict[str, Any], proto.Message],
    message: Optional[proto.Message] = None,
) -> proto.Message:
    """Parses a JSON-like object into a message.

    If the object is already a message, this will return the object as-is. If
    the object is a JSON Dict, this will parse and merge the object into the
    message.

    Args:
        obj (Union[dict[str, Any], proto.Message]):
            Required. The object to convert to a proto message.
        message (proto.Message):
            Optional. A protocol buffer message to merge the obj into. It
            defaults to Struct() if unspecified.

    Returns:
        proto.Message: The same message passed as argument.
    """
    if message is None:
        message = struct_pb2.Struct()
    if isinstance(obj, (proto.Message, struct_pb2.Struct)):
        return obj
    try:
        json_format.ParseDict(obj, message._pb)
    except AttributeError:
        json_format.ParseDict(obj, message)
    return message


def _upload_agent_engine(
    *,
    agent: _AgentEngineInterface,
    gcs_bucket: _StorageBucket,
    gcs_dir_name: str,
) -> None:
    """Uploads the agent engine to GCS."""
    cloudpickle = _import_cloudpickle_or_raise()
    blob = gcs_bucket.blob(f"{gcs_dir_name}/{_BLOB_FILENAME}")
    with blob.open("wb") as f:
        try:
            cloudpickle.dump(agent, f)
        except Exception as e:
            url = "https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop/custom#deployment-considerations"
            error_msg = f"Failed to serialize agent engine. Visit {url} for details."
            if "google._upb._message" in str(e) or "Descriptor" in str(e):
                error_msg += (
                    " This is often caused by protobuf objects (like Part, AgentCard) "
                    "being imported at the global module level. Please move these "
                    "imports inside the functions or methods where they are used. "
                    "Alternatively, you can import the entire module: "
                    "`from a2a import types`."
                )
            raise TypeError(error_msg) 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: _StorageBucket,
    gcs_dir_name: str,
) -> 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: _StorageBucket,
    gcs_dir_name: str,
) -> 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 _create_base64_encoded_tarball(
    *,
    source_packages: Sequence[str],
) -> str:
    """Creates a base64 encoded tarball from the source packages."""
    logger.info("Creating in-memory tarfile of source_packages")
    tar_fileobj = io.BytesIO()
    project_dir = os.path.realpath(os.getcwd())
    with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar:
        for file in source_packages:
            real_file_path = os.path.realpath(file)
            if real_file_path != project_dir and not real_file_path.startswith(
                project_dir + os.sep
            ):
                raise ValueError(
                    f"File path '{file}' is outside the project directory "
                    f"'{project_dir}'."
                )
            tar.add(file)
    tar_fileobj.seek(0)
    tarball_bytes = tar_fileobj.read()
    return base64.b64encode(tarball_bytes).decode("utf-8")


def _validate_packages_or_raise(
    *,
    packages: Sequence[str],
    build_options: Optional[Dict[str, Sequence[str]]] = None,
) -> Sequence[str]:
    """Tries to validates the packages."""
    packages = packages or []
    if build_options and _INSTALLATION_SUBDIR in build_options:
        _validate_installation_scripts_or_raise(
            script_paths=build_options[_INSTALLATION_SUBDIR],
            packages=packages,
        )
    for package in packages:
        if not os.path.exists(package):
            raise FileNotFoundError(f"Package specified but not found: {package=}")
    return packages


def _validate_installation_scripts_or_raise(
    *,
    script_paths: Sequence[str],
    packages: Sequence[str],
) -> None:
    """Validates the installation scripts' path explicitly provided by the user.

    Args:
        script_paths (Sequence[str]):
            Required. The paths to the installation scripts.
        packages (Sequence[str]):
            Required. The user-provided packages.

    Raises:
        ValueError: If a user-defined script is not under the expected
            subdirectory, or not in `packages`, or if a package is
            in the installation scripts subdirectory, but is not specified as an
            installation script.
    """
    for script_path in script_paths:
        if not script_path.startswith(_INSTALLATION_SUBDIR):
            logger.warning(
                f"User-defined installation script '{script_path}' is not in "
                f"the expected '{_INSTALLATION_SUBDIR}' subdirectory. "
                f"Ensure it is placed in '{_INSTALLATION_SUBDIR}' within your "
                f"'extra_packages' or 'source_packages'."
            )
            raise ValueError(
                f"Required installation script '{script_path}' "
                f"is not under '{_INSTALLATION_SUBDIR}'"
            )

        if script_path not in packages:
            logger.warning(
                f"User-defined installation script '{script_path}' is not in "
                f"'extra_packages' or 'source_packages'. Ensure it is added to "
                f"'extra_packages' or 'source_packages'."
            )
            raise ValueError(
                f"User-defined installation script '{script_path}' "
                f"does not exist in 'extra_packages' or 'source_packages'."
            )

    for package in packages:
        if package.startswith(_INSTALLATION_SUBDIR) and package not in script_paths:
            logger.warning(
                f"Package '{package}' is in the installation "
                "scripts subdirectory, but is not specified as an installation "
                "script in `build_options`. "
                "Ensure it is added to installation_scripts for "
                "automatic execution."
            )
            raise ValueError(
                f"Package '{package}' is in the installation "
                "scripts subdirectory, but is not specified as an installation "
                "script in `build_options`."
            )
    return


def _validate_staging_bucket_or_raise(*, staging_bucket: str) -> str:
    """Tries to validate the staging bucket."""
    if not staging_bucket:
        raise ValueError(
            "Please provide a `staging_bucket` in `client.agent_engines.create(...)`."
        )
    if not staging_bucket.startswith("gs://"):
        raise ValueError(f"{staging_bucket=} must start with `gs://`")
    return staging_bucket


def _validate_requirements_or_warn(
    *,
    obj: Any,
    requirements: List[str],
) -> List[str]:
    """Compiles the requirements into a list of requirements."""
    requirements = requirements.copy()
    try:
        current_requirements = _scan_requirements(obj=obj)
        logger.info(f"Identified the following requirements: {current_requirements}")
        constraints = _parse_constraints(constraints=requirements)
        missing_requirements = _compare_requirements(
            requirements=current_requirements,
            constraints=constraints,
        )
        for warning_type, warnings in missing_requirements["warnings"].items():
            if warnings:
                logger.warning(
                    f"The following requirements are {warning_type}: {warnings}"
                )
        for action_type, actions in missing_requirements["actions"].items():
            if actions and action_type == _ACTION_APPEND:
                for action in actions:  # type: ignore[attr-defined]
                    requirements.append(action)
                logger.info(f"The following requirements are appended: {actions}")
    except Exception as e:
        logger.warning(f"Failed to compile requirements: {e}")
    return requirements


def _validate_requirements_or_raise(
    *,
    agent: Any,
    requirements: Optional[Sequence[str]] = None,
) -> 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 = _validate_requirements_or_warn(
        obj=agent,
        requirements=requirements,
    )
    logger.info(f"The final list of requirements: {requirements}")
    return requirements


def _validate_agent_or_raise(
    *,
    agent: _AgentEngineInterface,
) -> _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: The agent to be validated.

    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, BaseAgent):
            logger.info("Deploying google.adk.agents.Agent as an application.")
            from agentplatform import agent_engines

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

    if not (
        is_queryable
        or is_async_queryable
        or is_stream_queryable
        or is_operation_registrable
        or is_bidi_stream_queryable
        or is_async_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, "query"))
        except ValueError as err:
            raise ValueError(
                "Invalid query signature. This might be due to a missing "
                "`self` argument in the agent.query method."
            ) from err

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

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

    if is_async_stream_queryable:
        try:
            inspect.signature(getattr(agent, "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.async_stream_query method."
            ) from err

    if is_bidi_stream_queryable:
        try:
            inspect.signature(getattr(agent, "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.bidi_stream_query method."
            ) from err

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

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


def _wrap_agent_operation(*, agent: Any, operation: str) -> Callable[..., Any]:
    """Wraps an agent operation into a method (works for all API modes)."""

    def _method(self, **kwargs) -> Any:  # type: ignore[no-untyped-def]
        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


def _wrap_query_operation(*, method_name: str) -> Callable[..., Any]:
    """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: genai_types.AgentEngine, **kwargs) -> Any:  # type: ignore[no-untyped-def]
        if not self.api_client:
            raise ValueError("api_client is not initialized.")
        if not self.api_resource:
            raise ValueError("api_resource is not initialized.")
        response = self.api_client._query(
            name=self.api_resource.name,
            config={
                "class_method": method_name,
                "input": kwargs,
                "include_all_fields": True,
            },
        )
        return response.output

    return _method


def _wrap_async_query_operation(
    *, method_name: str
) -> Callable[..., Coroutine[Any, Any, Any]]:
    """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: genai_types.AgentEngine, **kwargs: Any
    ) -> Union[Coroutine[Any, Any, Any], Any]:
        if not self.api_async_client:
            raise ValueError("api_async_client is not initialized.")
        if not self.api_resource:
            raise ValueError("api_resource is not initialized.")
        response = await self.api_async_client._query(
            name=self.api_resource.name,
            config={
                "class_method": method_name,
                "input": kwargs,
                "include_all_fields": True,
            },
        )
        return response.output

    return _method


def _wrap_stream_query_operation(*, method_name: str) -> Callable[..., Iterator[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: genai_types.AgentEngine, **kwargs) -> Iterator[Any]:  # type: ignore[no-untyped-def]
        if not self.api_client:
            raise ValueError("api_client is not initialized.")
        if not self.api_resource:
            raise ValueError("api_resource is not initialized.")
        for http_response in self.api_client._stream_query(
            name=self.api_resource.name,
            config={
                "class_method": method_name,
                "input": kwargs,
                "include_all_fields": True,
            },
        ):
            for line in _yield_parsed_json(http_response=http_response):
                if line is not None:
                    yield line

    return _method


def _wrap_async_stream_query_operation(
    *, method_name: str
) -> Callable[..., AsyncIterator[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: genai_types.AgentEngine, **kwargs) -> AsyncIterator[Any]:  # type: ignore[no-untyped-def]
        if not self.api_client:
            raise ValueError("api_client is not initialized.")
        if not self.api_resource:
            raise ValueError("api_resource is not initialized.")
        async for http_response in self.api_client._async_stream_query(
            name=self.api_resource.name,
            config={
                "class_method": method_name,
                "input": kwargs,
                "include_all_fields": True,
            },
        ):
            for line in _yield_parsed_json(http_response=http_response):
                if line is not None:
                    yield line

    return _method


def _wrap_a2a_operation(method_name: str, agent_card: str) -> Callable[..., list[Any]]:
    """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:  # type: ignore[no-untyped-def]
        """Wraps an Agent Engine method, creating a callable for A2A API."""
        if not self.api_client:
            raise ValueError("api_client is not initialized.")
        if not self.api_resource:
            raise ValueError("api_resource is not initialized.")
        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

        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.
        base_url = self.api_client._api_client._http_options.base_url.rstrip("/")
        api_version = self.api_client._api_client._http_options.api_version
        a2a_agent_card.url = f"{base_url}/{api_version}/{self.api_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 {self.api_client._api_client._credentials.token}"
                    )
                },
                timeout=(
                    self.api_client._api_client._http_options.timeout / 1000.0
                    if self.api_client._api_client._http_options.timeout
                    else None
                ),
            ),
        )
        factory = ClientFactory(config)
        client = factory.create(a2a_agent_card)

        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  # type: ignore[return-value]


def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]:
    """Converts the body of the HTTP Response message to JSON format.

    Args:
        http_response (google.genai.types.HttpResponse):
            Required. The httpbody body to be converted to JSON object(s).

    Yields:
        Any: A JSON object or line of the original body or None.
    """
    if not http_response.body:
        yield None
        return

    # Handle the case of multiple dictionaries delimited by newlines.
    for line in http_response.body.split("\n"):
        if line:
            try:
                line = json.loads(line)
            except Exception as e:
                logger.warning(f"failed to parse json: {line}. Exception: {e}")
            yield line


def _validate_resource_limits_or_raise(resource_limits: dict[str, str]) -> None:
    """Validates the resource limits.

    Checks that the resource limits are a dict with 'cpu' and 'memory' keys.
    Checks that the 'cpu' value is one of 1, 2, 4, 6, 8.
    Checks that the 'memory' value is a string ending with 'Gi'.
    Checks that the memory size is smaller than 32Gi.
    Checks that the memory size requires at least the specified number of CPUs.

    Args:
        resource_limits: The resource limits to be validated.

    Raises:
        TypeError: If the resource limits are not a dict.
        KeyError: If the resource limits do not contain 'cpu' and 'memory' keys.
        ValueError: If the 'cpu' value is not one of 1, 2, 4, 6, 8.
        ValueError: If the 'memory' value is not a string ending with 'Gi'.
        ValueError: If the memory size is too large.
        ValueError: If the memory size requires more CPUs than the specified
        'cpu' value.
    """
    if not isinstance(resource_limits, dict):
        raise TypeError(f"resource_limits must be a dict. Got {type(resource_limits)}")
    if "cpu" not in resource_limits or "memory" not in resource_limits:
        raise KeyError("resource_limits must contain 'cpu' and 'memory' keys.")

    cpu = int(resource_limits["cpu"])
    memory_str = resource_limits["memory"]

    if cpu not in [1, 2, 4, 6, 8]:
        raise ValueError(
            "resource_limits['cpu'] must be one of 1, 2, 4, 6, 8. Got" f" {cpu}"
        )

    if not isinstance(memory_str, str) or not memory_str.endswith("Gi"):
        raise ValueError(
            "resource_limits['memory'] must be a string ending with 'Gi'."
            f" Got {memory_str}"
        )

    try:
        memory_gb = int(memory_str[:-2])
    except ValueError:
        raise ValueError(
            f"Invalid memory value: {memory_str}. Must be an integer"
            " followed by 'Gi'."
        )

    # https://cloud.google.com/run/docs/configuring/memory-limits
    if memory_gb > 32:
        raise ValueError(
            f"Memory size of {memory_str} is too large. Must be smaller than 32Gi."
        )
    if memory_gb > 24:
        min_cpu = 8
    elif memory_gb > 16:
        min_cpu = 6
    elif memory_gb > 8:
        min_cpu = 4
    elif memory_gb > 4:
        min_cpu = 2
    else:
        min_cpu = 1

    if cpu < min_cpu:
        raise ValueError(
            f"Memory size of {memory_str} requires at least {min_cpu} CPUs."
            f" Got {cpu}"
        )


def _is_adk_agent(agent_engine: _AgentEngineInterface) -> bool:
    """Checks if the agent engine is an ADK agent.

    Args:
        agent_engine: The agent engine to check.

    Returns:
        True if the agent engine is an ADK agent, False otherwise.
    """

    from agentplatform.agent_engines.templates import adk

    return isinstance(agent_engine, adk.AdkApp)


def _add_telemetry_enablement_env(
    env_vars: Optional[Dict[str, Union[str, Any]]]
) -> Optional[Dict[str, Union[str, Any]]]:
    """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"
    )
    env_to_add = {GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY: "unspecified"}

    if env_vars is None:
        return env_to_add

    if not isinstance(env_vars, dict):
        raise TypeError(f"env_vars must be a dict, but got {type(env_vars)}.")

    if GOOGLE_CLOUD_AGENT_ENGINE_ENABLE_TELEMETRY in env_vars:
        return env_vars

    return env_vars | env_to_add


def _dataclass_to_dict_or_raise(obj: Any) -> Dict[str, Any]:
    """Converts a dataclass to a JSON dictionary."""
    if not dataclasses.is_dataclass(obj):
        raise TypeError(f"Object is not a dataclass: {obj}")
    return json.loads(json.dumps(dataclasses.asdict(obj)))


def _autogen_run_response_protocol_to_dict(
    obj: AutogenRunResponse,
) -> Dict[str, Any]:
    """Converts an AutogenRunResponse object into a JSON-serializable dictionary."""
    if hasattr(obj, "process"):
        obj.process()
    last_speaker = None
    if getattr(obj, "last_speaker", None) is not None:
        last_speaker = {
            "name": getattr(obj.last_speaker, "name", None),
            "description": getattr(obj.last_speaker, "description", None),
        }
    cost = None
    if getattr(obj, "cost", None) is not None:
        if hasattr(obj.cost, "model_dump_json"):
            cost = json.loads(obj.cost.model_dump_json())
        else:
            cost = str(obj.cost)
    result = {
        "summary": getattr(obj, "summary", None),
        "messages": list(getattr(obj, "messages", [])),
        "context_variables": getattr(obj, "context_variables", None),
        "last_speaker": last_speaker,
        "cost": cost,
    }
    return json.loads(json.dumps(result))


def to_json_serializable_autogen_object(
    obj: Union[
        AutogenChatResult,
        AutogenRunResponse,
    ],
) -> Dict[str, Any]:
    """Converts an Autogen object to a JSON serializable object."""
    if isinstance(obj, AutogenChatResult):
        return _dataclass_to_dict_or_raise(obj)
    return _autogen_run_response_protocol_to_dict(obj)


def _llama_index_response_to_dict(obj: LlamaIndexResponse) -> Any:
    response = {}
    if hasattr(obj, "response"):
        response["response"] = obj.response
    if hasattr(obj, "source_nodes"):
        response["source_nodes"] = [node.model_dump_json() for node in obj.source_nodes]
    if hasattr(obj, "metadata"):
        response["metadata"] = obj.metadata
    return json.loads(json.dumps(response))


def _llama_index_chat_response_to_dict(obj: LlamaIndexChatResponse) -> Any:
    return json.loads(obj.message.model_dump_json())


def _llama_index_base_model_to_dict(obj: LlamaIndexBaseModel) -> Any:
    return json.loads(obj.model_dump_json())


def to_json_serializable_llama_index_object(
    obj: Union[
        LlamaIndexResponse,
        LlamaIndexBaseModel,
        LlamaIndexChatResponse,
        Sequence[LlamaIndexBaseModel],
    ],
) -> Union[str, Dict[str, Any], Sequence[Union[str, Dict[str, Any]]]]:
    """Converts a LlamaIndexResponse to a JSON serializable object."""
    if isinstance(obj, LlamaIndexResponse):
        return _llama_index_response_to_dict(obj)
    if isinstance(obj, LlamaIndexChatResponse):
        return _llama_index_chat_response_to_dict(obj)
    if isinstance(obj, Sequence):
        seq_result = []
        for item in obj:
            if isinstance(item, LlamaIndexBaseModel):
                seq_result.append(_llama_index_base_model_to_dict(item))
                continue
            seq_result.append(str(item))
        return seq_result
    if isinstance(obj, LlamaIndexBaseModel):
        return _llama_index_base_model_to_dict(obj)
    return str(obj)


def is_noop_or_proxy_tracer_provider(tracer_provider) -> bool:
    """Returns True if the tracer_provider is Proxy or NoOp."""
    opentelemetry = _import_opentelemetry_or_warn()
    if not opentelemetry:
        return False
    ProxyTracerProvider = opentelemetry.trace.ProxyTracerProvider
    NoOpTracerProvider = opentelemetry.trace.NoOpTracerProvider
    return isinstance(tracer_provider, (NoOpTracerProvider, ProxyTracerProvider))


def dump_event_for_json(event: BaseModel) -> Dict[str, Any]:
    """Dumps an ADK event to a JSON-serializable dictionary."""
    return json.loads(event.model_dump_json(exclude_none=True))


def _import_opentelemetry_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the opentelemetry module."""
    try:
        import opentelemetry

        return opentelemetry
    except ImportError:
        logger.warning(
            "opentelemetry-api is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        )
    return None


def _import_opentelemetry_sdk_trace_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the opentelemetry.sdk.trace module."""
    try:
        import opentelemetry.sdk.trace

        return opentelemetry.sdk.trace
    except ImportError:
        logger.warning(
            "opentelemetry-sdk is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        )
    return None


def _import_cloud_trace_v2_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the google.cloud.trace_v2 module."""
    try:
        import google.cloud.trace_v2

        return google.cloud.trace_v2
    except ImportError:
        logger.warning(
            "google-cloud-trace is not installed. Please call "
            "'pip install google-cloud-aiplatform[agent_engines]'."
        )
    return None


def _import_cloud_trace_exporter_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the opentelemetry.exporter.cloud_trace module."""
    try:
        import opentelemetry.exporter.cloud_trace

        return opentelemetry.exporter.cloud_trace
    except ImportError:
        logger.warning(
            "opentelemetry-exporter-gcp-trace is not installed. Please "
            "call 'pip install google-cloud-aiplatform[agent_engines]'."
        )
    return None


def _import_openinference_langchain_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the openinference.instrumentation.langchain module."""
    try:
        import openinference.instrumentation.langchain

        return openinference.instrumentation.langchain
    except ImportError:
        logger.warning(
            "openinference-instrumentation-langchain is not installed. Please "
            "call 'pip install google-cloud-aiplatform[langchain]'."
        )
    return None


def _import_openinference_autogen_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the openinference.instrumentation.autogen module."""
    try:
        import openinference.instrumentation.autogen

        return openinference.instrumentation.autogen
    except ImportError:
        logger.warning(
            "openinference-instrumentation-autogen is not installed. Please "
            "call 'pip install google-cloud-aiplatform[ag2]'."
        )
    return None


def _import_openinference_llama_index_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the openinference.instrumentation.llama_index module."""
    try:
        import openinference.instrumentation.llama_index  # noqa:F401

        return openinference.instrumentation.llama_index
    except ImportError:
        logger.warning(
            "openinference-instrumentation-llama_index is not installed. Please "
            "call 'pip install google-cloud-aiplatform[llama_index]'."
        )
    return None


def _import_nest_asyncio_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the nest_asyncio module."""
    try:
        import nest_asyncio

        return nest_asyncio
    except ImportError:
        logger.warning(
            "nest_asyncio is not installed. Please call: `pip install nest-asyncio`"
        )
    return None


def _import_autogen_tools_or_warn() -> Optional[types.ModuleType]:
    """Tries to import the autogen.tools module."""
    try:
        from autogen import tools

        return tools
    except ImportError:
        logger.warning(
            "autogen.tools is not installed. Please "
            "call `pip install google-cloud-aiplatform[ag2]`."
        )
    return None