Skip to content

Commit 67fee04

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 e5f71de commit 67fee04

4 files changed

Lines changed: 396 additions & 83 deletions

File tree

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

Lines changed: 186 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -577,16 +577,16 @@ def register_operations(self) -> Dict[str, List[str]]:
577577
OperationRegistrableEngine().custom_method,
578578
schema_name=_TEST_CUSTOM_METHOD_NAME,
579579
)
580-
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
581-
_TEST_STANDARD_API_MODE
582-
)
580+
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA[
581+
_TEST_MODE_KEY_IN_SCHEMA
582+
] = _TEST_STANDARD_API_MODE
583583
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA = _agent_engines_utils._generate_schema(
584584
OperationRegistrableEngine().custom_async_method,
585585
schema_name=_TEST_CUSTOM_ASYNC_METHOD_NAME,
586586
)
587-
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
588-
_TEST_ASYNC_API_MODE
589-
)
587+
_TEST_AGENT_ENGINE_ASYNC_CUSTOM_METHOD_SCHEMA[
588+
_TEST_MODE_KEY_IN_SCHEMA
589+
] = _TEST_ASYNC_API_MODE
590590
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
591591
StreamQueryEngine().stream_query,
592592
schema_name=_TEST_DEFAULT_STREAM_METHOD_NAME,
@@ -596,41 +596,41 @@ def register_operations(self) -> Dict[str, List[str]]:
596596
OperationRegistrableEngine().custom_stream_method,
597597
schema_name=_TEST_CUSTOM_STREAM_METHOD_NAME,
598598
)
599-
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
600-
_TEST_STREAM_API_MODE
601-
)
599+
_TEST_AGENT_ENGINE_CUSTOM_STREAM_QUERY_SCHEMA[
600+
_TEST_MODE_KEY_IN_SCHEMA
601+
] = _TEST_STREAM_API_MODE
602602
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
603603
AsyncStreamQueryEngine().async_stream_query,
604604
schema_name=_TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME,
605605
)
606-
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
607-
_TEST_ASYNC_STREAM_API_MODE
608-
)
606+
_TEST_AGENT_ENGINE_ASYNC_STREAM_QUERY_SCHEMA[
607+
_TEST_MODE_KEY_IN_SCHEMA
608+
] = _TEST_ASYNC_STREAM_API_MODE
609609
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA = (
610610
_agent_engines_utils._generate_schema(
611611
OperationRegistrableEngine().custom_async_stream_method,
612612
schema_name=_TEST_CUSTOM_ASYNC_STREAM_METHOD_NAME,
613613
)
614614
)
615-
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
616-
_TEST_ASYNC_STREAM_API_MODE
617-
)
615+
_TEST_AGENT_ENGINE_CUSTOM_ASYNC_STREAM_QUERY_SCHEMA[
616+
_TEST_MODE_KEY_IN_SCHEMA
617+
] = _TEST_ASYNC_STREAM_API_MODE
618618
_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA = _agent_engines_utils._generate_schema(
619619
OperationRegistrableEngine().bidi_stream_query,
620620
schema_name=_TEST_DEFAULT_BIDI_STREAM_METHOD_NAME,
621621
)
622-
_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
623-
_TEST_BIDI_STREAM_API_MODE
624-
)
622+
_TEST_AGENT_ENGINE_BIDI_STREAM_QUERY_SCHEMA[
623+
_TEST_MODE_KEY_IN_SCHEMA
624+
] = _TEST_BIDI_STREAM_API_MODE
625625
_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA = (
626626
_agent_engines_utils._generate_schema(
627627
OperationRegistrableEngine().custom_bidi_stream_method,
628628
schema_name=_TEST_CUSTOM_BIDI_STREAM_METHOD_NAME,
629629
)
630630
)
631-
_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
632-
_TEST_BIDI_STREAM_API_MODE
633-
)
631+
_TEST_AGENT_ENGINE_CUSTOM_BIDI_STREAM_QUERY_SCHEMA[
632+
_TEST_MODE_KEY_IN_SCHEMA
633+
] = _TEST_BIDI_STREAM_API_MODE
634634
_TEST_OPERATION_REGISTRABLE_SCHEMAS = [
635635
_TEST_AGENT_ENGINE_QUERY_SCHEMA,
636636
_TEST_AGENT_ENGINE_CUSTOM_METHOD_SCHEMA,
@@ -657,9 +657,9 @@ def register_operations(self) -> Dict[str, List[str]]:
657657
MethodToBeUnregisteredEngine().method_to_be_unregistered,
658658
schema_name=_TEST_METHOD_TO_BE_UNREGISTERED_NAME,
659659
)
660-
_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA[_TEST_MODE_KEY_IN_SCHEMA] = (
661-
_TEST_STANDARD_API_MODE
662-
)
660+
_TEST_METHOD_TO_BE_UNREGISTERED_SCHEMA[
661+
_TEST_MODE_KEY_IN_SCHEMA
662+
] = _TEST_STANDARD_API_MODE
663663
_TEST_ASYNC_QUERY_SCHEMAS = [_TEST_AGENT_ENGINE_ASYNC_METHOD_SCHEMA]
664664
_TEST_STREAM_QUERY_SCHEMAS = [
665665
_TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA,
@@ -1113,6 +1113,63 @@ def test_create_agent_engine_config_with_developer_connect_source(self):
11131113
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
11141114
)
11151115

1116+
def test_create_agent_engine_config_with_agent_config_source_and_requirements_file(
1117+
self,
1118+
):
1119+
with tempfile.TemporaryDirectory() as tmpdir:
1120+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1121+
with open(requirements_file_path, "w") as f:
1122+
f.write("requests==2.0.0")
1123+
1124+
config = self.client.agent_engines._create_config(
1125+
mode="create",
1126+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1127+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1128+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1129+
agent_framework=_TEST_AGENT_FRAMEWORK,
1130+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1131+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1132+
agent_config_source={"adk_config": {"json_config": {}}},
1133+
requirements_file=requirements_file_path,
1134+
)
1135+
1136+
assert config["spec"]["source_code_spec"] == {
1137+
"agent_config_source": {"adk_config": {"json_config": {}}},
1138+
"python_spec": {
1139+
"version": _TEST_PYTHON_VERSION_OVERRIDE,
1140+
"requirements_file": requirements_file_path,
1141+
},
1142+
}
1143+
1144+
def test_create_agent_engine_config_with_agent_config_source_and_entrypoint_module_warns(
1145+
self, caplog
1146+
):
1147+
caplog.set_level(logging.WARNING, logger="vertexai_genai.agentengines")
1148+
1149+
config = self.client.agent_engines._create_config(
1150+
mode="create",
1151+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1152+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1153+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1154+
agent_framework=_TEST_AGENT_FRAMEWORK,
1155+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1156+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1157+
agent_config_source={"adk_config": {"json_config": {}}},
1158+
entrypoint_module="some_module",
1159+
)
1160+
1161+
assert (
1162+
"`entrypoint_module` and `entrypoint_object` are ignored when"
1163+
in caplog.text
1164+
)
1165+
assert config["spec"]["source_code_spec"] == {
1166+
"agent_config_source": {"adk_config": {"json_config": {}}},
1167+
"python_spec": {
1168+
"version": _TEST_PYTHON_VERSION_OVERRIDE,
1169+
},
1170+
}
1171+
# entrypoint_module is NOT in python_spec
1172+
11161173
@mock.patch.object(
11171174
_agent_engines_utils,
11181175
"_create_base64_encoded_tarball",
@@ -1146,6 +1203,106 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises(
11461203
)
11471204
assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
11481205

1206+
@mock.patch.object(
1207+
_agent_engines_utils,
1208+
"_create_base64_encoded_tarball",
1209+
return_value="test_tarball",
1210+
)
1211+
def test_create_agent_engine_config_with_agent_config_source_and_image_spec_raises(
1212+
self, mock_create_base64_encoded_tarball
1213+
):
1214+
with tempfile.TemporaryDirectory() as tmpdir:
1215+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1216+
with open(test_file_path, "w") as f:
1217+
f.write("test content")
1218+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1219+
with open(requirements_file_path, "w") as f:
1220+
f.write("requests==2.0.0")
1221+
1222+
with pytest.raises(ValueError) as excinfo:
1223+
self.client.agent_engines._create_config(
1224+
mode="create",
1225+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1226+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1227+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1228+
agent_framework=_TEST_AGENT_FRAMEWORK,
1229+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1230+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1231+
image_spec={},
1232+
agent_config_source={"adk_config": {"json_config": {}}},
1233+
)
1234+
assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
1235+
1236+
def test_create_agent_engine_config_with_agent_config_source(self):
1237+
config = self.client.agent_engines._create_config(
1238+
mode="create",
1239+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1240+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1241+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1242+
agent_framework=_TEST_AGENT_FRAMEWORK,
1243+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1244+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1245+
agent_config_source={"adk_config": {"json_config": {}}},
1246+
)
1247+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1248+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1249+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1250+
assert config["spec"]["source_code_spec"] == {
1251+
"agent_config_source": {"adk_config": {"json_config": {}}},
1252+
"python_spec": {"version": _TEST_PYTHON_VERSION_OVERRIDE},
1253+
}
1254+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1255+
assert (
1256+
config["spec"]["identity_type"]
1257+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1258+
)
1259+
1260+
@mock.patch.object(
1261+
_agent_engines_utils,
1262+
"_create_base64_encoded_tarball",
1263+
return_value="test_tarball",
1264+
)
1265+
def test_create_agent_engine_config_with_source_packages_and_agent_config_source(
1266+
self, mock_create_base64_encoded_tarball
1267+
):
1268+
with tempfile.TemporaryDirectory() as tmpdir:
1269+
test_file_path = os.path.join(tmpdir, "test_file.txt")
1270+
with open(test_file_path, "w") as f:
1271+
f.write("test content")
1272+
requirements_file_path = os.path.join(tmpdir, "requirements.txt")
1273+
with open(requirements_file_path, "w") as f:
1274+
f.write("requests==2.0.0")
1275+
1276+
config = self.client.agent_engines._create_config(
1277+
mode="create",
1278+
display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
1279+
description=_TEST_AGENT_ENGINE_DESCRIPTION,
1280+
source_packages=[test_file_path],
1281+
class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
1282+
agent_framework=_TEST_AGENT_FRAMEWORK,
1283+
identity_type=_TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT,
1284+
python_version=_TEST_PYTHON_VERSION_OVERRIDE,
1285+
agent_config_source={"adk_config": {"json_config": {}}},
1286+
)
1287+
assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
1288+
assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
1289+
assert config["spec"]["agent_framework"] == _TEST_AGENT_FRAMEWORK
1290+
assert config["spec"]["source_code_spec"] == {
1291+
"agent_config_source": {
1292+
"adk_config": {"json_config": {}},
1293+
"inline_source": {"source_archive": "test_tarball"},
1294+
},
1295+
"python_spec": {"version": _TEST_PYTHON_VERSION_OVERRIDE},
1296+
}
1297+
assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
1298+
mock_create_base64_encoded_tarball.assert_called_once_with(
1299+
source_packages=[test_file_path]
1300+
)
1301+
assert (
1302+
config["spec"]["identity_type"]
1303+
== _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
1304+
)
1305+
11491306
@mock.patch.object(
11501307
_agent_engines_utils,
11511308
"_create_base64_encoded_tarball",
@@ -1916,6 +2073,7 @@ def test_create_agent_engine_with_env_vars_dict(
19162073
python_version=None,
19172074
build_options=None,
19182075
image_spec=None,
2076+
agent_config_source=None,
19192077
)
19202078
request_mock.assert_called_with(
19212079
"post",
@@ -2018,6 +2176,7 @@ def test_create_agent_engine_with_custom_service_account(
20182176
python_version=None,
20192177
build_options=None,
20202178
image_spec=None,
2179+
agent_config_source=None,
20212180
)
20222181
request_mock.assert_called_with(
20232182
"post",
@@ -2119,6 +2278,7 @@ def test_create_agent_engine_with_experimental_mode(
21192278
python_version=None,
21202279
build_options=None,
21212280
image_spec=None,
2281+
agent_config_source=None,
21222282
)
21232283
request_mock.assert_called_with(
21242284
"post",
@@ -2289,6 +2449,7 @@ def test_create_agent_engine_with_class_methods(
22892449
python_version=None,
22902450
build_options=None,
22912451
image_spec=None,
2452+
agent_config_source=None,
22922453
)
22932454
request_mock.assert_called_with(
22942455
"post",
@@ -2385,6 +2546,7 @@ def test_create_agent_engine_with_agent_framework(
23852546
python_version=None,
23862547
build_options=None,
23872548
image_spec=None,
2549+
agent_config_source=None,
23882550
)
23892551
request_mock.assert_called_with(
23902552
"post",

0 commit comments

Comments
 (0)