Skip to content

Commit 856b33c

Browse files
shawn-yang-googlecopybara-github
shawn-yang-google
authored and
copybara-github
committed
feat: add support for agent_config_source (ADK) in AgentEngines
PiperOrigin-RevId: 875355220
1 parent 5705565 commit 856b33c

4 files changed

Lines changed: 375 additions & 4 deletions

File tree

‎tests/unit/vertexai/genai/test_agent_engines.py‎

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1146,6 +1146,232 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises(
11461146
)
11471147
assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
11481148

1149+
@mock.patch.object(
1150+
_agent_engines_utils,
1151+
"_create_base64_encoded_tarball",
1152+
return_value="test_tarball",
1153+
)
1154+
def test_create_agent_engine_config_with_agent_config_source_and_image_spec_raises(
1155+
self, mock_create_base64_encoded_tarball
1156+
):
1157+
with tempfile.TemporaryDirectory() as tmpdir:
1158+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1159+
with open(test_file_path, "w") as f:
1160+
f.write("test content")
1161+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1162+
with open(requirements_file_path, "w") as f:
1163+
f.write("requests==2.0.0")
1164+
1165+
with pytest.raises(ValueError) as excinfo:
1166+
self.client.agent_engines._create_config(
1167+
mode="create",
1168+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1169+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1170+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1171+
agent_framework=_TEST_AGENT_FRAMEWORK,
1172+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1173+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1174+
image_spec={},
1175+
agent_config_source={"adk_config": {"json_config": {}}},
1176+
)
1177+
assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
1178+
1179+
def test_create_agent_engine_config_with_agent_config_source(self):
1180+
config = self.client.agent_engines._create_config(
1181+
mode="create",
1182+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1183+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1184+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1185+
agent_framework=_TEST_AGENT_FRAMEWORK,
1186+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1187+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1188+
agent_config_source={"adk_config": {"json_config": {}}},
1189+
)
1190+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1191+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1192+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1193+
assert config["spec"]["source_code_spec"] == {
1194+
"agent_config_source": {"adk_config": {"json_config": {}}},
1195+
}
1196+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1197+
assert (
1198+
config["spec"]["identity_type"]
1199+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1200+
)
1201+
1202+
@mock.patch.object(
1203+
_agent_engines_utils,
1204+
"_create_base64_encoded_tarball",
1205+
return_value="test_tarball",
1206+
)
1207+
def test_create_agent_engine_config_with_source_packages_and_agent_config_source(
1208+
self, mock_create_base64_encoded_tarball
1209+
):
1210+
with tempfile.TemporaryDirectory() as tmpdir:
1211+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1212+
with open(test_file_path, "w") as f:
1213+
f.write("test content")
1214+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1215+
with open(requirements_file_path, "w") as f:
1216+
f.write("requests==2.0.0")
1217+
1218+
config = self.client.agent_engines._create_config(
1219+
mode="create",
1220+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1221+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1222+
source_packages=[test_file_path],
1223+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1224+
agent_framework=_TEST_AGENT_FRAMEWORK,
1225+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1226+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1227+
agent_config_source={"adk_config": {"json_config": {}}},
1228+
)
1229+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1230+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1231+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1232+
assert config["spec"]["source_code_spec"] == {
1233+
"agent_config_source": {
1234+
"adk_config": {"json_config": {}},
1235+
"inline_source": {"source_archive": "test_tarball"},
1236+
},
1237+
}
1238+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1239+
mock_create_base64_encoded_tarball.assert_called_once_with(
1240+
source_packages=[test_file_path]
1241+
)
1242+
assert (
1243+
config["spec"]["identity_type"]
1244+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1245+
)
1246+
1247+
@mock.patch.object(
1248+
_agent_engines_utils,
1249+
"_create_base64_encoded_tarball",
1250+
return_value="test_tarball",
1251+
)
1252+
def test_create_agent_engine_config_with_agent_config_source_and_image_spec_raises(
1253+
self, mock_create_base64_encoded_tarball
1254+
):
1255+
with tempfile.TemporaryDirectory() as tmpdir:
1256+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1257+
with open(test_file_path, "w") as f:
1258+
f.write("test content")
1259+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1260+
with open(requirements_file_path, "w") as f:
1261+
f.write("requests==2.0.0")
1262+
1263+
with pytest.raises(ValueError) as excinfo:
1264+
self.client.agent_engines._create_config(
1265+
mode="create",
1266+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1267+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1268+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1269+
agent_framework=_TEST_AGENT_FRAMEWORK,
1270+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1271+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1272+
image_spec={},
1273+
agent_config_source={"adk_config": {"json_config": {}}},
1274+
)
1275+
assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
1276+
1277+
def test_create_agent_engine_config_with_agent_config_source(self):
1278+
config = self.client.agent_engines._create_config(
1279+
mode="create",
1280+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1281+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1282+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1283+
agent_framework=_TEST_AGENT_FRAMEWORK,
1284+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1285+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1286+
agent_config_source={"adk_config": {"json_config": {}}},
1287+
)
1288+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1289+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1290+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1291+
assert config["spec"]["source_code_spec"] == {
1292+
"agent_config_source": {"adk_config": {"json_config": {}}},
1293+
}
1294+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1295+
assert (
1296+
config["spec"]["identity_type"]
1297+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1298+
)
1299+
1300+
@mock.patch.object(
1301+
_agent_engines_utils,
1302+
"_create_base64_encoded_tarball",
1303+
return_value="test_tarball",
1304+
)
1305+
def test_create_agent_engine_config_with_source_packages_and_agent_config_source(
1306+
self, mock_create_base64_encoded_tarball
1307+
):
1308+
with tempfile.TemporaryDirectory() as tmpdir:
1309+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1310+
with open(test_file_path, "w") as f:
1311+
f.write("test content")
1312+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1313+
with open(requirements_file_path, "w") as f:
1314+
f.write("requests==2.0.0")
1315+
1316+
config = self.client.agent_engines._create_config(
1317+
mode="create",
1318+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1319+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1320+
source_packages=[test_file_path],
1321+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1322+
agent_framework=_TEST_AGENT_FRAMEWORK,
1323+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1324+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1325+
agent_config_source={"adk_config": {"json_config": {}}},
1326+
)
1327+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1328+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1329+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1330+
assert config["spec"]["source_code_spec"] == {
1331+
"agent_config_source": {
1332+
"adk_config": {"json_config": {}},
1333+
"inline_source": {"source_archive": "test_tarball"},
1334+
},
1335+
}
1336+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1337+
mock_create_base64_encoded_tarball.assert_called_once_with(
1338+
source_packages=[test_file_path]
1339+
)
1340+
assert (
1341+
config["spec"]["identity_type"]
1342+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1343+
)
1344+
1345+
@mock.patch.object(
1346+
_agent_engines_utils,
1347+
"_create_base64_encoded_tarball",
1348+
return_value="test_tarball",
1349+
)
1350+
def test_create_agent_engine_config_with_agent_config_source_and_image_spec_raises(
1351+
self, mock_create_base64_encoded_tarball
1352+
):
1353+
with tempfile.TemporaryDirectory() as tmpdir:
1354+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1355+
with open(test_file_path, "w") as f:
1356+
f.write("test content")
1357+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1358+
with open(requirements_file_path, "w") as f:
1359+
f.write("requests==2.0.0")
1360+
1361+
with pytest.raises(ValueError) as excinfo:
1362+
self.client.agent_engines._create_config(
1363+
mode="create",
1364+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1365+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1366+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1367+
agent_framework=_TEST_AGENT_FRAMEWORK,
1368+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1369+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1370+
image_spec={},
1371+
agent_config_source={"adk_config": {"json_config": {}}},
1372+
)
1373+
assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
1374+
11491375
@mock.patch.object(
11501376
_agent_engines_utils,
11511377
"_create_base64_encoded_tarball",
@@ -1916,6 +2142,7 @@ def test_create_agent_engine_with_env_vars_dict(
19162142
python_version=None,
19172143
build_options=None,
19182144
image_spec=None,
2145+
agent_config_source=None,
19192146
)
19202147
request_mock.assert_called_with(
19212148
"post",
@@ -2018,6 +2245,7 @@ def test_create_agent_engine_with_custom_service_account(
20182245
python_version=None,
20192246
build_options=None,
20202247
image_spec=None,
2248+
agent_config_source=None,
20212249
)
20222250
request_mock.assert_called_with(
20232251
"post",
@@ -2119,6 +2347,7 @@ def test_create_agent_engine_with_experimental_mode(
21192347
python_version=None,
21202348
build_options=None,
21212349
image_spec=None,
2350+
agent_config_source=None,
21222351
)
21232352
request_mock.assert_called_with(
21242353
"post",
@@ -2289,6 +2518,7 @@ def test_create_agent_engine_with_class_methods(
22892518
python_version=None,
22902519
build_options=None,
22912520
image_spec=None,
2521+
agent_config_source=None,
22922522
)
22932523
request_mock.assert_called_with(
22942524
"post",
@@ -2385,6 +2615,7 @@ def test_create_agent_engine_with_agent_framework(
23852615
python_version=None,
23862616
build_options=None,
23872617
image_spec=None,
2618+
agent_config_source=None,
23882619
)
23892620
request_mock.assert_called_with(
23902621
"post",

‎vertexai/_genai/agent_engines.py‎

Lines changed: 51 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -786,7 +786,12 @@ def _list_pager(
786786
def _is_lightweight_creation(
787787
self, agent: Any, config: types.AgentEngineConfig
788788
) -> bool:
789-
if agent or config.source_packages or config.developer_connect_source:
789+
if (
790+
agent
791+
or config.source_packages
792+
or config.developer_connect_source
793+
or config.agent_config_source
794+
):
790795
return False
791796
return True
792797

@@ -975,6 +980,7 @@ def create(
975980
python_version=config.python_version,
976981
build_options=config.build_options,
977982
image_spec=config.image_spec,
983+
agent_config_source=config.agent_config_source,
978984
)
979985
operation = self._create(config=api_config)
980986
reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
@@ -1034,10 +1040,13 @@ def _set_source_code_spec(
10341040
image_spec: Optional[
10351041
types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
10361042
] = None,
1043+
agent_config_source: Optional[
1044+
types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict
1045+
] = None,
10371046
) -> None:
10381047
"""Sets source_code_spec for agent engine inside the `spec`."""
10391048
source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict()
1040-
if source_packages:
1049+
if source_packages and not agent_config_source:
10411050
source_packages = _agent_engines_utils._validate_packages_or_raise(
10421051
packages=source_packages,
10431052
build_options=build_options,
@@ -1053,7 +1062,7 @@ def _set_source_code_spec(
10531062
source_code_spec["developer_connect_source"] = {
10541063
"config": developer_connect_source
10551064
}
1056-
else:
1065+
elif not agent_config_source:
10571066
raise ValueError(
10581067
"Please specify one of `source_packages` or `developer_connect_source`."
10591068
)
@@ -1078,11 +1087,40 @@ def _set_source_code_spec(
10781087
"`entrypoint_object`, or `requirements_file`, as they are "
10791088
"mutually exclusive."
10801089
)
1090+
if agent_config_source:
1091+
raise ValueError(
1092+
"`image_spec` cannot be specified alongside `agent_config_source`, "
1093+
"as they are mutually exclusive."
1094+
)
10811095
update_masks.append("spec.source_code_spec.image_spec")
10821096
source_code_spec["image_spec"] = image_spec
10831097
spec["source_code_spec"] = source_code_spec
10841098
return
10851099

1100+
if agent_config_source is not None:
1101+
if entrypoint_module or entrypoint_object or requirements_file:
1102+
raise ValueError(
1103+
"`agent_config_source` cannot be specified alongside `entrypoint_module`, "
1104+
"`entrypoint_object`, or `requirements_file`, as they are "
1105+
"mutually exclusive."
1106+
)
1107+
if source_packages:
1108+
source_packages = _agent_engines_utils._validate_packages_or_raise(
1109+
packages=source_packages,
1110+
build_options=build_options,
1111+
)
1112+
if "inline_source" not in agent_config_source:
1113+
agent_config_source["inline_source"] = {}
1114+
agent_config_source["inline_source"]["source_archive"] = (
1115+
_agent_engines_utils._create_base64_encoded_tarball(
1116+
source_packages=source_packages
1117+
)
1118+
)
1119+
update_masks.append("spec.source_code_spec.agent_config_source")
1120+
source_code_spec["agent_config_source"] = agent_config_source
1121+
spec["source_code_spec"] = source_code_spec
1122+
return
1123+
10861124
update_masks.append("spec.source_code_spec.python_spec.version")
10871125
python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = {
10881126
"version": sys_version,
@@ -1235,6 +1273,9 @@ def _create_config(
12351273
image_spec: Optional[
12361274
types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
12371275
] = None,
1276+
agent_config_source: Optional[
1277+
types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict
1278+
] = None,
12381279
) -> types.UpdateAgentEngineConfigDict:
12391280
import sys
12401281

@@ -1307,7 +1348,12 @@ def _create_config(
13071348
sys_version=sys_version,
13081349
build_options=build_options,
13091350
)
1310-
elif source_packages or developer_connect_source:
1351+
elif (
1352+
source_packages
1353+
or developer_connect_source
1354+
or image_spec
1355+
or agent_config_source
1356+
):
13111357
agent_engine_spec = {}
13121358
self._set_source_code_spec(
13131359
spec=agent_engine_spec,
@@ -1321,6 +1367,7 @@ def _create_config(
13211367
sys_version=sys_version,
13221368
build_options=build_options,
13231369
image_spec=image_spec,
1370+
agent_config_source=agent_config_source,
13241371
)
13251372

13261373
is_deployment_spec_updated = (

0 commit comments

Comments
 (0)