feat: add support for agent_config_source (ADK) in AgentEngines

PiperOrigin-RevId: 875355220

Author: shawn-yang-google

tests/unit/vertexai/genai/test_agent_engines.py

@@ -1113,6 +1113,63 @@ def test_create_agent_engine_config_with_developer_connect_source(self):
                 == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
             )
 
+    def test_create_agent_engine_config_with_agent_config_source_and_requirements_file(
+        self,
+    ):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            requirements_file_path = os.path.join(tmpdir, "requirements.txt")
+            with open(requirements_file_path, "w") as f:
+                f.write("requests==2.0.0")
+
+            config = self.client.agent_engines._create_config(
+                mode="create",
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                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,
+                agent_config_source={"adk_config": {"json_config": {}}},
+                requirements_file=requirements_file_path,
+            )
+
+            assert config["spec"]["source_code_spec"] == {
+                "agent_config_source": {"adk_config": {"json_config": {}}},
+                "python_spec": {
+                    "version": _TEST_PYTHON_VERSION_OVERRIDE,
+                    "requirements_file": requirements_file_path,
+                },
+            }
+
+    def test_create_agent_engine_config_with_agent_config_source_and_entrypoint_module_warns(
+        self, caplog
+    ):
+        caplog.set_level(logging.WARNING, logger="vertexai_genai.agentengines")
+
+        config = self.client.agent_engines._create_config(
+            mode="create",
+            display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+            description=_TEST_AGENT_ENGINE_DESCRIPTION,
+            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,
+            agent_config_source={"adk_config": {"json_config": {}}},
+            entrypoint_module="some_module",
+        )
+
+        assert (
+            "`entrypoint_module` and `entrypoint_object` are ignored when"
+            in caplog.text
+        )
+        assert config["spec"]["source_code_spec"] == {
+            "agent_config_source": {"adk_config": {"json_config": {}}},
+            "python_spec": {
+                "version": _TEST_PYTHON_VERSION_OVERRIDE,
+            },
+        }
+        # entrypoint_module is NOT in python_spec
+
     @mock.patch.object(
         _agent_engines_utils,
         "_create_base64_encoded_tarball",
@@ -1146,6 +1203,106 @@ def test_create_agent_engine_config_with_source_packages_and_image_spec_raises(
                 )
             assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
 
+    @mock.patch.object(
+        _agent_engines_utils,
+        "_create_base64_encoded_tarball",
+        return_value="test_tarball",
+    )
+    def test_create_agent_engine_config_with_agent_config_source_and_image_spec_raises(
+        self, mock_create_base64_encoded_tarball
+    ):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_file_path = os.path.join(tmpdir, "test_file.txt")
+            with open(test_file_path, "w") as f:
+                f.write("test content")
+            requirements_file_path = os.path.join(tmpdir, "requirements.txt")
+            with open(requirements_file_path, "w") as f:
+                f.write("requests==2.0.0")
+
+            with pytest.raises(ValueError) as excinfo:
+                self.client.agent_engines._create_config(
+                    mode="create",
+                    display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                    description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                    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,
+                    image_spec={},
+                    agent_config_source={"adk_config": {"json_config": {}}},
+                )
+            assert "`image_spec` cannot be specified alongside" in str(excinfo.value)
+
+    def test_create_agent_engine_config_with_agent_config_source(self):
+        config = self.client.agent_engines._create_config(
+            mode="create",
+            display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+            description=_TEST_AGENT_ENGINE_DESCRIPTION,
+            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,
+            agent_config_source={"adk_config": {"json_config": {}}},
+        )
+        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"] == {
+            "agent_config_source": {"adk_config": {"json_config": {}}},
+            "python_spec": {"version": _TEST_PYTHON_VERSION_OVERRIDE},
+        }
+        assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
+        assert (
+            config["spec"]["identity_type"]
+            == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+        )
+
+    @mock.patch.object(
+        _agent_engines_utils,
+        "_create_base64_encoded_tarball",
+        return_value="test_tarball",
+    )
+    def test_create_agent_engine_config_with_source_packages_and_agent_config_source(
+        self, mock_create_base64_encoded_tarball
+    ):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test_file_path = os.path.join(tmpdir, "test_file.txt")
+            with open(test_file_path, "w") as f:
+                f.write("test content")
+            requirements_file_path = os.path.join(tmpdir, "requirements.txt")
+            with open(requirements_file_path, "w") as f:
+                f.write("requests==2.0.0")
+
+            config = self.client.agent_engines._create_config(
+                mode="create",
+                display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                source_packages=[test_file_path],
+                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,
+                agent_config_source={"adk_config": {"json_config": {}}},
+            )
+            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"] == {
+                "agent_config_source": {
+                    "adk_config": {"json_config": {}},
+                    "inline_source": {"source_archive": "test_tarball"},
+                },
+                "python_spec": {"version": _TEST_PYTHON_VERSION_OVERRIDE},
+            }
+            assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
+            mock_create_base64_encoded_tarball.assert_called_once_with(
+                source_packages=[test_file_path]
+            )
+            assert (
+                config["spec"]["identity_type"]
+                == _TEST_AGENT_ENGINE_IDENTITY_TYPE_SERVICE_ACCOUNT
+            )
+
     @mock.patch.object(
         _agent_engines_utils,
         "_create_base64_encoded_tarball",
@@ -1916,6 +2073,7 @@ def test_create_agent_engine_with_env_vars_dict(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                agent_config_source=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2018,6 +2176,7 @@ def test_create_agent_engine_with_custom_service_account(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                agent_config_source=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2119,6 +2278,7 @@ def test_create_agent_engine_with_experimental_mode(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                agent_config_source=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2289,6 +2449,7 @@ def test_create_agent_engine_with_class_methods(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                agent_config_source=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -2385,6 +2546,7 @@ def test_create_agent_engine_with_agent_framework(
                 python_version=None,
                 build_options=None,
                 image_spec=None,
+                agent_config_source=None,
             )
             request_mock.assert_called_with(
                 "post",

vertexai/_genai/agent_engines.py

@@ -786,7 +786,12 @@ def _list_pager(
     def _is_lightweight_creation(
         self, agent: Any, config: types.AgentEngineConfig
     ) -> bool:
-        if agent or config.source_packages or config.developer_connect_source:
+        if (
+            agent
+            or config.source_packages
+            or config.developer_connect_source
+            or config.agent_config_source
+        ):
             return False
         return True
 
@@ -937,6 +942,9 @@ def create(
             developer_connect_source = json.loads(
                 developer_connect_source.model_dump_json()
             )
+        agent_config_source = config.agent_config_source
+        if agent_config_source is not None:
+            agent_config_source = json.loads(agent_config_source.model_dump_json())
         if agent and agent_engine:
             raise ValueError("Please specify only one of `agent` or `agent_engine`.")
         elif agent_engine:
@@ -975,6 +983,7 @@ def create(
             python_version=config.python_version,
             build_options=config.build_options,
             image_spec=config.image_spec,
+            agent_config_source=agent_config_source,
         )
         operation = self._create(config=api_config)
         reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(
@@ -1034,10 +1043,13 @@ def _set_source_code_spec(
         image_spec: Optional[
             types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
         ] = None,
+        agent_config_source: Optional[
+            types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict
+        ] = None,
     ) -> None:
         """Sets source_code_spec for agent engine inside the `spec`."""
         source_code_spec = types.ReasoningEngineSpecSourceCodeSpecDict()
-        if source_packages:
+        if source_packages and not agent_config_source:
             source_packages = _agent_engines_utils._validate_packages_or_raise(
                 packages=source_packages,
                 build_options=build_options,
@@ -1053,13 +1065,15 @@ def _set_source_code_spec(
             source_code_spec["developer_connect_source"] = {
                 "config": developer_connect_source
             }
-        else:
+        elif not agent_config_source:
             raise ValueError(
-                "Please specify one of `source_packages` or `developer_connect_source`."
+                "Please specify one of `source_packages`, `developer_connect_source`, "
+                "or `agent_config_source`."
             )
         if class_methods is None:
             raise ValueError(
-                "`class_methods` must be specified if `source_packages` or `developer_connect_source` is specified."
+                "`class_methods` must be specified if `source_packages`, "
+                "`developer_connect_source`, or `agent_config_source` is specified."
             )
         update_masks.append("spec.class_methods")
         class_methods_spec_list = (
@@ -1078,6 +1092,11 @@ def _set_source_code_spec(
                     "`entrypoint_object`, or `requirements_file`, as they are "
                     "mutually exclusive."
                 )
+            if agent_config_source:
+                raise ValueError(
+                    "`image_spec` cannot be specified alongside `agent_config_source`, "
+                    "as they are mutually exclusive."
+                )
             update_masks.append("spec.source_code_spec.image_spec")
             source_code_spec["image_spec"] = image_spec
             spec["source_code_spec"] = source_code_spec
@@ -1087,6 +1106,38 @@ def _set_source_code_spec(
         python_spec: types.ReasoningEngineSpecSourceCodeSpecPythonSpecDict = {
             "version": sys_version,
         }
+        if agent_config_source is not None:
+            if entrypoint_module or entrypoint_object:
+                logger.warning(
+                    "`entrypoint_module` and `entrypoint_object` are ignored when "
+                    "`agent_config_source` is specified, as they are pre-defined."
+                )
+            if source_packages:
+                source_packages = _agent_engines_utils._validate_packages_or_raise(
+                    packages=source_packages,
+                    build_options=build_options,
+                )
+                update_masks.append(
+                    "spec.source_code_spec.agent_config_source.inline_source.source_archive"
+                )
+                agent_config_source["inline_source"] = {  # type: ignore[typeddict-item]
+                    "source_archive": _agent_engines_utils._create_base64_encoded_tarball(
+                        source_packages=source_packages
+                    )
+                }
+            update_masks.append("spec.source_code_spec.agent_config_source")
+            source_code_spec["agent_config_source"] = agent_config_source
+
+            if requirements_file is not None:
+                update_masks.append(
+                    "spec.source_code_spec.python_spec.requirements_file"
+                )
+                python_spec["requirements_file"] = requirements_file
+            source_code_spec["python_spec"] = python_spec
+
+            spec["source_code_spec"] = source_code_spec
+            return
+
         if not entrypoint_module:
             raise ValueError(
                 "`entrypoint_module` must be specified if `source_packages` or `developer_connect_source` is specified."
@@ -1235,6 +1286,9 @@ def _create_config(
         image_spec: Optional[
             types.ReasoningEngineSpecSourceCodeSpecImageSpecDict
         ] = None,
+        agent_config_source: Optional[
+            types.ReasoningEngineSpecSourceCodeSpecAgentConfigSourceDict
+        ] = None,
     ) -> types.UpdateAgentEngineConfigDict:
         import sys
 
@@ -1307,7 +1361,12 @@ def _create_config(
                 sys_version=sys_version,
                 build_options=build_options,
             )
-        elif source_packages or developer_connect_source:
+        elif (
+            source_packages
+            or developer_connect_source
+            or image_spec
+            or agent_config_source
+        ):
             agent_engine_spec = {}
             self._set_source_code_spec(
                 spec=agent_engine_spec,
@@ -1321,6 +1380,7 @@ def _create_config(
                 sys_version=sys_version,
                 build_options=build_options,
                 image_spec=image_spec,
+                agent_config_source=agent_config_source,
             )
 
         is_deployment_spec_updated = (
@@ -1336,7 +1396,8 @@ def _create_config(
                 "To update `env_vars`, `psc_interface_config`, `min_instances`, "
                 "`max_instances`, `resource_limits`, or `container_concurrency`, "
                 "you must also provide the `agent` variable or the source code "
-                "options (`source_packages` or `developer_connect_source`)."
+                "options (`source_packages`, `developer_connect_source` or "
+                "`agent_config_source`)."
             )
 
         if agent_engine_spec is not None:
@@ -1598,6 +1659,9 @@ def update(
             developer_connect_source = json.loads(
                 developer_connect_source.model_dump_json()
             )
+        agent_config_source = config.agent_config_source
+        if agent_config_source is not None:
+            agent_config_source = json.loads(agent_config_source.model_dump_json())
         if agent and agent_engine:
             raise ValueError("Please specify only one of `agent` or `agent_engine`.")
         elif agent_engine:
@@ -1633,6 +1697,7 @@ def update(
             agent_framework=config.agent_framework,
             python_version=config.python_version,
             build_options=config.build_options,
+            agent_config_source=agent_config_source,
         )
         operation = self._update(name=name, config=api_config)
         reasoning_engine_id = _agent_engines_utils._get_reasoning_engine_id(