Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 12 additions & 3 deletions tests/unit/vertexai/genai/test_agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,7 @@ def register_operations(self) -> Dict[str, List[str]]:
)
_TEST_AGENT_ENGINE_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = _TEST_STANDARD_API_MODE
_TEST_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
_TEST_PYTHON_VERSION_OVERRIDE = "3.11"
_TEST_AGENT_ENGINE_FRAMEWORK = _agent_engines_utils._DEFAULT_AGENT_FRAMEWORK
_TEST_AGENT_ENGINE_CLASS_METHOD_1 = {
"description": "Runs the engine.",
Expand Down Expand Up @@ -969,12 +970,13 @@ def test_create_agent_engine_config_full(self, mock_prepare):
resource_limits=_TEST_AGENT_ENGINE_RESOURCE_LIMITS,
container_concurrency=_TEST_AGENT_ENGINE_CONTAINER_CONCURRENCY,
encryption_spec=_TEST_AGENT_ENGINE_ENCRYPTION_SPEC,
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
assert config["spec"]["agent_framework"] == "custom"
assert config["spec"]["package_spec"] == {
"python_version": _TEST_PYTHON_VERSION,
"python_version": _TEST_PYTHON_VERSION_OVERRIDE,
"pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
"dependency_files_gcs_uri": _TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI,
"requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
Expand Down Expand Up @@ -1037,14 +1039,15 @@ def test_create_agent_engine_config_with_source_packages(
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
agent_framework=_TEST_AGENT_FRAMEWORK,
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
assert config["spec"]["source_code_spec"] == {
"inline_source": {"source_archive": "test_tarball"},
"python_spec": {
"version": _TEST_PYTHON_VERSION,
"version": _TEST_PYTHON_VERSION_OVERRIDE,
"entrypoint_module": "main",
"entrypoint_object": "app",
"requirements_file": requirements_file_path,
Expand Down Expand Up @@ -1073,12 +1076,13 @@ def test_update_agent_engine_config_full(self, mock_prepare):
env_vars=_TEST_AGENT_ENGINE_ENV_VARS_INPUT,
service_account=_TEST_AGENT_ENGINE_CUSTOM_SERVICE_ACCOUNT,
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
)
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
assert config["spec"]["agent_framework"] == "custom"
assert config["spec"]["package_spec"] == {
"python_version": _TEST_PYTHON_VERSION,
"python_version": _TEST_PYTHON_VERSION_OVERRIDE,
"pickle_object_gcs_uri": _TEST_AGENT_ENGINE_GCS_URI,
"dependency_files_gcs_uri": _TEST_AGENT_ENGINE_DEPENDENCY_FILES_GCS_URI,
"requirements_gcs_uri": _TEST_AGENT_ENGINE_REQUIREMENTS_GCS_URI,
Expand Down Expand Up @@ -1593,6 +1597,7 @@ def test_create_agent_engine_with_env_vars_dict(
entrypoint_object=None,
requirements_file=None,
agent_framework=None,
python_version=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -1685,6 +1690,7 @@ def test_create_agent_engine_with_custom_service_account(
entrypoint_object=None,
requirements_file=None,
agent_framework=None,
python_version=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -1776,6 +1782,7 @@ def test_create_agent_engine_with_experimental_mode(
entrypoint_object=None,
requirements_file=None,
agent_framework=None,
python_version=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -1930,6 +1937,7 @@ def test_create_agent_engine_with_class_methods(
entrypoint_object=None,
requirements_file=None,
agent_framework=None,
python_version=None,
)
request_mock.assert_called_with(
"post",
Expand Down Expand Up @@ -2016,6 +2024,7 @@ def test_create_agent_engine_with_agent_framework(
requirements_file=None,
agent_framework=_TEST_AGENT_FRAMEWORK,
identity_type=None,
python_version=None,
)
request_mock.assert_called_with(
"post",
Expand Down
14 changes: 13 additions & 1 deletion vertexai/_genai/agent_engines.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,9 @@ def _CreateAgentEngineConfig_to_vertex(
if getv(from_object, ["agent_framework"]) is not None:
setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"]))

if getv(from_object, ["python_version"]) is not None:
setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"]))

return to_object


Expand Down Expand Up @@ -291,6 +294,9 @@ def _UpdateAgentEngineConfig_to_vertex(
if getv(from_object, ["agent_framework"]) is not None:
setv(parent_object, ["agentFramework"], getv(from_object, ["agent_framework"]))

if getv(from_object, ["python_version"]) is not None:
setv(parent_object, ["pythonVersion"], getv(from_object, ["python_version"]))

if getv(from_object, ["update_mask"]) is not None:
setv(
parent_object, ["_query", "updateMask"], getv(from_object, ["update_mask"])
Expand Down Expand Up @@ -931,6 +937,7 @@ def create(
entrypoint_object=config.entrypoint_object,
requirements_file=config.requirements_file,
agent_framework=config.agent_framework,
python_version=config.python_version,
)
operation = self._create(config=api_config)
# TODO: Use a more specific link.
Expand Down Expand Up @@ -996,6 +1003,7 @@ def _create_config(
entrypoint_object: Optional[str] = None,
requirements_file: Optional[str] = None,
agent_framework: Optional[str] = None,
python_version: Optional[str] = None,
) -> types.UpdateAgentEngineConfigDict:
import sys

Expand Down Expand Up @@ -1027,7 +1035,10 @@ def _create_config(
if agent_framework == "google-adk":
env_vars = _agent_engines_utils._add_telemetry_enablement_env(env_vars)

sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
if python_version:
sys_version = python_version
else:
sys_version = f"{sys.version_info.major}.{sys.version_info.minor}"
agent_engine_spec = None
if agent is not None:
if source_packages is not None:
Expand Down Expand Up @@ -1453,6 +1464,7 @@ def update(
entrypoint_object=config.entrypoint_object,
requirements_file=config.requirements_file,
agent_framework=config.agent_framework,
python_version=config.python_version,
)
operation = self._update(name=name, config=api_config)
logger.info(
Expand Down
39 changes: 39 additions & 0 deletions vertexai/_genai/types/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -5457,6 +5457,13 @@ class CreateAgentEngineConfig(_common.BaseModel):
- If `source_packages` is specified, the agent framework will
default to "custom".""",
)
python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] = Field(
default=None,
description="""The Python version to be used for the Agent Engine.
If not specified, it will use the current Python version of the environment.
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
""",
)


class CreateAgentEngineConfigDict(TypedDict, total=False):
Expand Down Expand Up @@ -5567,6 +5574,12 @@ class CreateAgentEngineConfigDict(TypedDict, total=False):
- If `source_packages` is specified, the agent framework will
default to "custom"."""

python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]]
"""The Python version to be used for the Agent Engine.
If not specified, it will use the current Python version of the environment.
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
"""


CreateAgentEngineConfigOrDict = Union[
CreateAgentEngineConfig, CreateAgentEngineConfigDict
Expand Down Expand Up @@ -6183,6 +6196,13 @@ class UpdateAgentEngineConfig(_common.BaseModel):
- If `source_packages` is specified, the agent framework will
default to "custom".""",
)
python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] = Field(
default=None,
description="""The Python version to be used for the Agent Engine.
If not specified, it will use the current Python version of the environment.
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
""",
)
update_mask: Optional[str] = Field(
default=None,
description="""The update mask to apply. For the `FieldMask` definition, see
Expand Down Expand Up @@ -6298,6 +6318,12 @@ class UpdateAgentEngineConfigDict(TypedDict, total=False):
- If `source_packages` is specified, the agent framework will
default to "custom"."""

python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]]
"""The Python version to be used for the Agent Engine.
If not specified, it will use the current Python version of the environment.
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
"""

update_mask: Optional[str]
"""The update mask to apply. For the `FieldMask` definition, see
https://protobuf.dev/reference/protobuf/google.protobuf/#field-mask."""
Expand Down Expand Up @@ -13037,6 +13063,13 @@ class AgentEngineConfig(_common.BaseModel):
- If `source_packages` is specified, the agent framework will
default to "custom".""",
)
python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]] = Field(
default=None,
description="""The Python version to be used for the Agent Engine.
If not specified, it will use the current Python version of the environment.
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
""",
)


class AgentEngineConfigDict(TypedDict, total=False):
Expand Down Expand Up @@ -13179,6 +13212,12 @@ class AgentEngineConfigDict(TypedDict, total=False):
- If `source_packages` is specified, the agent framework will
default to "custom"."""

python_version: Optional[Literal["3.9", "3.10", "3.11", "3.12", "3.13"]]
"""The Python version to be used for the Agent Engine.
If not specified, it will use the current Python version of the environment.
Supported versions: "3.9", "3.10", "3.11", "3.12", "3.13".
"""


AgentEngineConfigOrDict = Union[AgentEngineConfig, AgentEngineConfigDict]

Expand Down
Loading