Content-Length: 1006943 | pFad | https://github.com/googleapis/python-aiplatform/commit/691c1f6156f2d66db369c2e5656c0e733b23a8f4

3F feat: Add support for env_vars parameter when creating or updating Ag… · googleapis/python-aiplatform@691c1f6 · GitHub
Skip to content

Commit 691c1f6

Browse files
yeesiancopybara-github
yeesian
authored and
copybara-github
committed
feat: Add support for env_vars parameter when creating or updating Agent Engine.
PiperOrigin-RevId: 747644161
1 parent 418607a commit 691c1f6

File tree

3 files changed

+424
-25
lines changed

3 files changed

+424
-25
lines changed

‎tests/unit/vertex_langchain/test_agent_engines.py

Lines changed: 232 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,8 @@ def register_operations(self) -> Dict[str, List[str]]:
274274
"""
275275
_TEST_METHOD_TO_BE_UNREGISTERED_NAME = "method_to_be_unregistered"
276276
_TEST_QUERY_PROMPT = "Find the first fibonacci number greater than 999"
277+
_TEST_AGENT_ENGINE_ENV_KEY = "GOOGLE_CLOUD_AGENT_ENGINE_ENV"
278+
_TEST_AGENT_ENGINE_ENV_VALUE = "test_env_value"
277279
_TEST_AGENT_ENGINE_GCS_URI = "{}/{}/{}".format(
278280
_TEST_STAGING_BUCKET,
279281
_TEST_GCS_DIR_NAME,
@@ -673,6 +675,8 @@ class TestAgentEngine:
673675
def setup_method(self):
674676
importlib.reload(initializer)
675677
importlib.reload(aiplatform)
678+
importlib.reload(os)
679+
os.environ[_TEST_AGENT_ENGINE_ENV_KEY] = _TEST_AGENT_ENGINE_ENV_VALUE
676680
aiplatform.init(
677681
project=_TEST_PROJECT,
678682
location=_TEST_LOCATION,
@@ -801,6 +805,119 @@ def test_create_agent_engine_requirements_from_file(
801805
retry=_TEST_RETRY,
802806
)
803807

808+
def test_create_agent_engine_with_env_vars_dict(
809+
self,
810+
create_agent_engine_mock,
811+
cloud_storage_create_bucket_mock,
812+
tarfile_open_mock,
813+
cloudpickle_dump_mock,
814+
cloudpickle_load_mock,
815+
importlib_metadata_version_mock,
816+
get_agent_engine_mock,
817+
get_gca_resource_mock,
818+
):
819+
agent_engines.create(
820+
self.test_agent,
821+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
822+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
823+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
824+
env_vars={
825+
"TEST_ENV_VAR": "TEST_ENV_VAR_VALUE",
826+
"TEST_ENV_VAR_2": "TEST_ENV_VAR_VALUE_2",
827+
"TEST_SECRET_ENV_VAR": {
828+
"secret": "TEST_SECRET_NAME_1",
829+
"version": "TEST_SECRET_VERSION_1",
830+
},
831+
"TEST_SECRET_ENV_VAR_2": types.SecretRef(
832+
secret="TEST_SECRET_NAME_2",
833+
version="TEST_SECRET_VERSION_2",
834+
),
835+
},
836+
)
837+
test_spec = types.ReasoningEngineSpec(
838+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
839+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
840+
env=[
841+
types.EnvVar(name="TEST_ENV_VAR", value="TEST_ENV_VAR_VALUE"),
842+
types.EnvVar(name="TEST_ENV_VAR_2", value="TEST_ENV_VAR_VALUE_2"),
843+
],
844+
secret_env=[
845+
types.SecretEnvVar(
846+
name="TEST_SECRET_ENV_VAR",
847+
secret_ref={
848+
"secret": "TEST_SECRET_NAME_1",
849+
"version": "TEST_SECRET_VERSION_1",
850+
},
851+
),
852+
types.SecretEnvVar(
853+
name="TEST_SECRET_ENV_VAR_2",
854+
secret_ref=types.SecretRef(
855+
secret="TEST_SECRET_NAME_2",
856+
version="TEST_SECRET_VERSION_2",
857+
),
858+
),
859+
],
860+
),
861+
)
862+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
863+
create_agent_engine_mock.assert_called_with(
864+
parent=_TEST_PARENT,
865+
reasoning_engine=types.ReasoningEngine(
866+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
867+
spec=test_spec,
868+
),
869+
)
870+
get_agent_engine_mock.assert_called_with(
871+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
872+
retry=_TEST_RETRY,
873+
)
874+
875+
def test_create_agent_engine_with_env_vars_list(
876+
self,
877+
create_agent_engine_mock,
878+
cloud_storage_create_bucket_mock,
879+
tarfile_open_mock,
880+
cloudpickle_dump_mock,
881+
cloudpickle_load_mock,
882+
importlib_metadata_version_mock,
883+
get_agent_engine_mock,
884+
get_gca_resource_mock,
885+
):
886+
agent_engines.create(
887+
self.test_agent,
888+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
889+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
890+
extra_packages=[_TEST_AGENT_ENGINE_EXTRA_PACKAGE_PATH],
891+
env_vars=[_TEST_AGENT_ENGINE_ENV_KEY, _TEST_AGENT_ENGINE_ENV_KEY],
892+
)
893+
test_spec = types.ReasoningEngineSpec(
894+
package_spec=_TEST_AGENT_ENGINE_PACKAGE_SPEC,
895+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
896+
env=[
897+
types.EnvVar(
898+
name=_TEST_AGENT_ENGINE_ENV_KEY,
899+
value=_TEST_AGENT_ENGINE_ENV_VALUE,
900+
),
901+
types.EnvVar(
902+
name=_TEST_AGENT_ENGINE_ENV_KEY,
903+
value=_TEST_AGENT_ENGINE_ENV_VALUE,
904+
),
905+
],
906+
),
907+
)
908+
test_spec.class_methods.append(_TEST_AGENT_ENGINE_QUERY_SCHEMA)
909+
create_agent_engine_mock.assert_called_with(
910+
parent=_TEST_PARENT,
911+
reasoning_engine=types.ReasoningEngine(
912+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
913+
spec=test_spec,
914+
),
915+
)
916+
get_agent_engine_mock.assert_called_with(
917+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
918+
retry=_TEST_RETRY,
919+
)
920+
804921
# pytest does not allow absl.testing.parameterized.named_parameters.
805922
@pytest.mark.parametrize(
806923
"test_case_name, test_kwargs, want_request",
@@ -923,6 +1040,48 @@ def test_create_agent_engine_requirements_from_file(
9231040
update_mask=field_mask_pb2.FieldMask(paths=["description"]),
9241041
),
9251042
),
1043+
(
1044+
"Update the environment variables",
1045+
{
1046+
"env_vars": {
1047+
_TEST_AGENT_ENGINE_ENV_KEY: _TEST_AGENT_ENGINE_ENV_VALUE,
1048+
"TEST_SECRET_ENV_VAR": {
1049+
"secret": "TEST_SECRET_NAME",
1050+
"version": "TEST_SECRET_VERSION",
1051+
},
1052+
},
1053+
},
1054+
types.reasoning_engine_service.UpdateReasoningEngineRequest(
1055+
reasoning_engine=types.ReasoningEngine(
1056+
name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
1057+
spec=types.ReasoningEngineSpec(
1058+
deployment_spec=types.ReasoningEngineSpec.DeploymentSpec(
1059+
env=[
1060+
types.EnvVar(
1061+
name=_TEST_AGENT_ENGINE_ENV_KEY,
1062+
value=_TEST_AGENT_ENGINE_ENV_VALUE,
1063+
),
1064+
],
1065+
secret_env=[
1066+
types.SecretEnvVar(
1067+
name="TEST_SECRET_ENV_VAR",
1068+
secret_ref=types.SecretRef(
1069+
secret="TEST_SECRET_NAME",
1070+
version="TEST_SECRET_VERSION",
1071+
),
1072+
),
1073+
],
1074+
),
1075+
),
1076+
),
1077+
update_mask=field_mask_pb2.FieldMask(
1078+
paths=[
1079+
"spec.deployment_spec.env",
1080+
"spec.deployment_spec.secret_env",
1081+
],
1082+
),
1083+
),
1084+
),
9261085
],
9271086
)
9281087
def test_update_agent_engine(
@@ -1832,6 +1991,78 @@ def test_create_agent_engine_with_invalid_register_operations_method(
18321991
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
18331992
)
18341993

1994+
def test_create_agent_engine_with_invalid_secret_ref_env_var(
1995+
self,
1996+
create_agent_engine_mock,
1997+
cloud_storage_create_bucket_mock,
1998+
tarfile_open_mock,
1999+
cloudpickle_dump_mock,
2000+
cloudpickle_load_mock,
2001+
importlib_metadata_version_mock,
2002+
get_agent_engine_mock,
2003+
):
2004+
with pytest.raises(ValueError, match="Failed to convert to secret ref"):
2005+
agent_engines.create(
2006+
self.test_agent,
2007+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
2008+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
2009+
env_vars={
2010+
"TEST_ENV_VAR": {
2011+
"name": "TEST_SECRET_NAME", # "name" should be "secret"
2012+
"version": "TEST_SECRET_VERSION",
2013+
},
2014+
},
2015+
)
2016+
2017+
def test_create_agent_engine_with_unknown_env_var(
2018+
self,
2019+
create_agent_engine_mock,
2020+
cloud_storage_create_bucket_mock,
2021+
tarfile_open_mock,
2022+
cloudpickle_dump_mock,
2023+
cloudpickle_load_mock,
2024+
importlib_metadata_version_mock,
2025+
get_agent_engine_mock,
2026+
):
2027+
with pytest.raises(ValueError, match="Env var not found in os.environ"):
2028+
agent_engines.create(
2029+
self.test_agent,
2030+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
2031+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
2032+
# Assumption: "UNKNOWN_TEST_ENV_VAR" not in os.environ
2033+
env_vars=["UNKNOWN_TEST_ENV_VAR"],
2034+
)
2035+
2036+
def test_create_agent_engine_with_invalid_type_env_var(
2037+
self,
2038+
create_agent_engine_mock,
2039+
cloud_storage_create_bucket_mock,
2040+
tarfile_open_mock,
2041+
cloudpickle_dump_mock,
2042+
cloudpickle_load_mock,
2043+
importlib_metadata_version_mock,
2044+
get_agent_engine_mock,
2045+
):
2046+
with pytest.raises(TypeError, match="Unknown value type in env_vars"):
2047+
agent_engines.create(
2048+
self.test_agent,
2049+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
2050+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
2051+
env_vars={
2052+
"TEST_ENV_VAR": 0.01, # should be a string or dict or SecretRef
2053+
},
2054+
)
2055+
with pytest.raises(TypeError, match="env_vars must be a list or a dict"):
2056+
agent_engines.create(
2057+
self.test_agent,
2058+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
2059+
requirements=_TEST_AGENT_ENGINE_REQUIREMENTS,
2060+
env_vars=types.SecretRef( # should be a list or dict
2061+
secret="TEST_SECRET_NAME",
2062+
version="TEST_SECRET_VERSION",
2063+
),
2064+
)
2065+
18352066
def test_update_agent_engine_unspecified_staging_bucket(
18362067
self,
18372068
update_agent_engine_mock,
@@ -1963,7 +2194,7 @@ def test_update_agent_engine_with_no_updates(
19632194
ValueError,
19642195
match=(
19652196
"At least one of `agent_engine`, `requirements`, "
1966-
"`extra_packages`, `display_name`, or `description` "
2197+
"`extra_packages`, `display_name`, `description`, or `env_vars` "
19672198
"must be specified."
19682199
),
19692200
):

‎vertexai/agent_engines/__init__.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,12 @@
1414
#
1515
"""Classes and functions for working with agent engines."""
1616

17-
from typing import Iterable, Optional, Sequence, Union
17+
from typing import Dict, Iterable, Optional, Sequence, Union
1818

1919
from google.cloud.aiplatform import base
2020
from google.cloud.aiplatform import initializer
2121
from google.cloud.aiplatform import utils as aip_utils
22-
from google.cloud.aiplatform_v1.types import (
23-
reasoning_engine_service as aip_types,
24-
)
22+
from google.cloud.aiplatform_v1 import types as aip_types
2523

2624
# We just want to re-export certain classes
2725
# pylint: disable=g-multiple-import,g-importing-member
@@ -66,6 +64,9 @@ def create(
6664
description: Optional[str] = None,
6765
gcs_dir_name: Optional[str] = None,
6866
extra_packages: Optional[Sequence[str]] = None,
67+
env_vars: Optional[
68+
Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
69+
] = None,
6970
) -> AgentEngine:
7071
"""Creates a new Agent Engine.
7172
@@ -121,6 +122,12 @@ def create(
121122
use for staging the artifacts needed.
122123
extra_packages (Sequence[str]):
123124
Optional. The set of extra user-provided packages (if any).
125+
env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]):
126+
Optional. The environment variables to be set when running the
127+
Agent Engine. If it is a list of strings, each string should be
128+
a valid key to `os.environ`. If it is a dictionary, the keys are
129+
the environment variable names, and the values are the
130+
corresponding values.
124131
125132
Returns:
126133
AgentEngine: The Agent Engine that was created.
@@ -142,6 +149,7 @@ def create(
142149
description=description,
143150
gcs_dir_name=gcs_dir_name,
144151
extra_packages=extra_packages,
152+
env_vars=env_vars,
145153
)
146154

147155

@@ -223,6 +231,9 @@ def update(
223231
description: Optional[str] = None,
224232
gcs_dir_name: Optional[str] = None,
225233
extra_packages: Optional[Sequence[str]] = None,
234+
env_vars: Optional[
235+
Union[Sequence[str], Dict[str, Union[str, aip_types.SecretRef]]]
236+
] = None,
226237
) -> "AgentEngine":
227238
"""Updates an existing Agent Engine.
228239
@@ -260,6 +271,12 @@ def update(
260271
it is not specified, the existing extra packages will be used.
261272
If it is set to an empty list, the existing extra packages will
262273
be removed.
274+
env_vars (Union[Sequence[str], Dict[str, Union[str, SecretRef]]]):
275+
Optional. The environment variables to be set when running the
276+
Agent Engine. If it is a list of strings, each string should be
277+
a valid key to `os.environ`. If it is a dictionary, the keys are
278+
the environment variable names, and the values are the
279+
corresponding values.
263280
264281
Returns:
265282
AgentEngine: The Agent Engine that was updated.
@@ -283,6 +300,7 @@ def update(
283300
description=description,
284301
gcs_dir_name=gcs_dir_name,
285302
extra_packages=extra_packages,
303+
env_vars=env_vars,
286304
)
287305

288306

0 commit comments

Comments
 (0)








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: https://github.com/googleapis/python-aiplatform/commit/691c1f6156f2d66db369c2e5656c0e733b23a8f4

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy