feat: Support Inline Source Deployment in Agent Engine

PiperOrigin-RevId: 817827304

Author: shawn-yang-google

tests/unit/vertexai/genai/replays/conftest.py

@@ -22,6 +22,7 @@
 from vertexai._genai import (
     client as vertexai_genai_client_module,
 )
+from vertexai._genai import _agent_engines_utils
 from google.cloud import storage, bigquery
 from google.genai import _replay_api_client
 from google.genai import client as google_genai_client_module
@@ -122,6 +123,16 @@ def replays_prefix():
     return "test"
 
 
+@pytest.fixture
+def mock_agent_engine_create_base64_encoded_tarball():
+    """Mocks the _create_base64_encoded_tarball function."""
+    with mock.patch.object(
+        _agent_engines_utils, "_create_base64_encoded_tarball"
+    ) as mock_create_base64_encoded_tarball:
+        mock_create_base64_encoded_tarball.return_value = "H4sIAAAAAAAAA-3UvWrDMBAHcM9-CpEpGRLkD8VQ6JOUElT7LFxkydEHxG9f2V1CKXSyu_x_i6TjJN2gk6N7HByNZIK_hEfINsCTa82zilcNTyMvRSPKao2vBM8KwZu6vJZXITJepGyRMb5FMT9FH6RjLHsM0mpr1CyN-i1vcsMo3aycjdMede0kV9YqTedW29id5TBpGXrrxjep0pO4kVGDIf-e_3edsI1APtxG20VNl2ne5o6_-r-oRer_Ypk2dd0s_Z82oP_3kLdaes-ensFLzpKOenaP5OajJ92fvoMLRyE6ww7LjrTwkzWeDnm-nmA_PqkN7PX5vOMJnwcAAAAAAAAAAAAAAAAAAADAdr4AI-kzQQAoAAA="
+        yield mock_create_base64_encoded_tarball
+
+
 def _get_replay_id(use_vertex: bool, replays_prefix: str) -> str:
     test_name_ending = os.environ.get("PYTEST_CURRENT_TEST").split("::")[-1]
     test_name = test_name_ending.split(" ")[0].split("[")[0] + "." + "vertex"

tests/unit/vertexai/genai/replays/test_create_agent_engine.py

@@ -19,6 +19,10 @@
 from tests.unit.vertexai.genai.replays import pytest_helper
 from vertexai._genai import types
 
+_TEST_CLASS_METHODS = [
+    {"name": "query", "api_mode": ""},
+]
+
 
 def test_create_config_lightweight(client):
     agent_display_name = "test-display-name"
@@ -108,6 +112,36 @@ def test_create_with_context_spec(client):
     client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
 
 
+def test_create_with_source_packages(
+    client,
+    mock_agent_engine_create_base64_encoded_tarball,
+):
+    """Tests creating an agent engine with source packages."""
+    with mock_agent_engine_create_base64_encoded_tarball:
+        agent_engine = client.agent_engines.create(
+            config={
+                "display_name": "test-agent-engine-source-packages",
+                "source_packages": [
+                    "test_module.py",
+                    "requirements.txt",
+                ],
+                "entrypoint_module": "test_module",
+                "entrypoint_object": "test_object",
+                "class_methods": _TEST_CLASS_METHODS,
+                "http_options": {
+                    "base_url": "https://us-west1-aiplatform.googleapis.com",
+                    "api_version": "v1beta1",
+                },
+            },
+        )
+        assert (
+            agent_engine.api_resource.display_name
+            == "test-agent-engine-source-packages"
+        )
+        # Clean up resources.
+        client.agent_engines.delete(name=agent_engine.api_resource.name, force=True)
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

tests/unit/vertexai/genai/test_agent_engines.py

@@ -13,11 +13,14 @@
 # limitations under the License.
 #
 import asyncio
+import base64
 import importlib
+import io
 import json
 import logging
 import os
 import sys
+import tarfile
 import tempfile
 from typing import Any, AsyncIterable, Dict, Iterable, List
 from unittest import mock
@@ -901,6 +904,48 @@ def test_create_agent_engine_config_full(self, mock_prepare):
             == _TEST_AGENT_ENGINE_CUSTOM_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(
+        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],
+                entrypoint_module="main",
+                entrypoint_object="app",
+                requirements_file=requirements_file_path,
+                class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+            )
+            assert config["display_name"] == _TEST_AGENT_ENGINE_DISPLAY_NAME
+            assert config["description"] == _TEST_AGENT_ENGINE_DESCRIPTION
+            assert config["spec"]["source_code_spec"] == {
+                "inline_source": {"source_archive": "test_tarball"},
+                "python_spec": {
+                    "version": _TEST_PYTHON_VERSION,
+                    "entrypoint_module": "main",
+                    "entrypoint_object": "app",
+                    "requirements_file": requirements_file_path,
+                },
+            }
+            assert config["spec"]["class_methods"] == _TEST_AGENT_ENGINE_CLASS_METHODS
+            mock_create_base64_encoded_tarball.assert_called_once_with(
+                source_packages=[test_file_path]
+            )
+
     @mock.patch.object(_agent_engines_utils, "_prepare")
     def test_update_agent_engine_config_full(self, mock_prepare):
         config = self.client.agent_engines._create_config(
@@ -951,10 +996,10 @@ def test_update_agent_engine_config_full(self, mock_prepare):
                 "spec.package_spec.pickle_object_gcs_uri",
                 "spec.package_spec.dependency_files_gcs_uri",
                 "spec.package_spec.requirements_gcs_uri",
+                "spec.class_methods",
                 "spec.deployment_spec.env",
                 "spec.deployment_spec.secret_env",
                 "spec.service_account",
-                "spec.class_methods",
                 "spec.agent_framework",
             ]
         )
@@ -1170,6 +1215,45 @@ def test_to_parsed_json(self, obj, expected):
         for got, want in zip(_agent_engines_utils._yield_parsed_json(obj), expected):
             assert got == want
 
+    def test_create_base64_encoded_tarball(self):
+        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")
+
+            origin_dir = os.getcwd()
+            try:
+                os.chdir(tmpdir)
+                encoded_tarball = _agent_engines_utils._create_base64_encoded_tarball(
+                    source_packages=["test_file.txt"]
+                )
+            finally:
+                os.chdir(origin_dir)
+
+            decoded_tarball = base64.b64decode(encoded_tarball)
+            with tarfile.open(fileobj=io.BytesIO(decoded_tarball), mode="r:gz") as tar:
+                names = tar.getnames()
+                assert "test_file.txt" in names
+
+    def test_create_base64_encoded_tarball_outside_project_dir_raises(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            project_dir = os.path.join(tmpdir, "project")
+            os.makedirs(project_dir)
+            sibling_path = os.path.join(tmpdir, "sibling.txt")
+            with open(sibling_path, "w") as f:
+                f.write("test content")
+
+            origin_dir = os.getcwd()
+            try:
+                os.chdir(project_dir)
+                with pytest.raises(ValueError) as excinfo:
+                    _agent_engines_utils._create_base64_encoded_tarball(
+                        source_packages=["../sibling.txt"]
+                    )
+                assert "is outside the project directory" in str(excinfo.value)
+            finally:
+                os.chdir(origin_dir)
+
 
 @pytest.mark.usefixtures("google_auth_mock")
 class TestAgentEngine:
@@ -1365,6 +1449,10 @@ def test_create_agent_engine_with_env_vars_dict(
                 agent_server_mode=None,
                 labels=None,
                 class_methods=None,
+                source_packages=None,
+                entrypoint_module=None,
+                entrypoint_object=None,
+                requirements_file=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1447,6 +1535,10 @@ def test_create_agent_engine_with_custom_service_account(
                 labels=None,
                 agent_server_mode=None,
                 class_methods=None,
+                source_packages=None,
+                entrypoint_module=None,
+                entrypoint_object=None,
+                requirements_file=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1531,6 +1623,10 @@ def test_create_agent_engine_with_experimental_mode(
                 labels=None,
                 agent_server_mode=_genai_types.AgentServerMode.EXPERIMENTAL,
                 class_methods=None,
+                source_packages=None,
+                entrypoint_module=None,
+                entrypoint_object=None,
+                requirements_file=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1553,6 +1649,72 @@ def test_create_agent_engine_with_experimental_mode(
                 None,
             )
 
+    @mock.patch.object(
+        _agent_engines_utils,
+        "_create_base64_encoded_tarball",
+        return_value="test_tarball",
+    )
+    @mock.patch.object(_agent_engines_utils, "_await_operation")
+    def test_create_agent_engine_with_source_packages(
+        self,
+        mock_await_operation,
+        mock_create_base64_encoded_tarball,
+    ):
+        mock_await_operation.return_value = _genai_types.AgentEngineOperation(
+            response=_genai_types.ReasoningEngine(
+                name=_TEST_AGENT_ENGINE_RESOURCE_NAME,
+                spec=_TEST_AGENT_ENGINE_SPEC,
+            )
+        )
+        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 mock.patch.object(
+                self.client.agent_engines._api_client, "request"
+            ) as request_mock:
+                request_mock.return_value = genai_types.HttpResponse(body="")
+                self.client.agent_engines.create(
+                    config=_genai_types.AgentEngineConfig(
+                        display_name=_TEST_AGENT_ENGINE_DISPLAY_NAME,
+                        description=_TEST_AGENT_ENGINE_DESCRIPTION,
+                        source_packages=[test_file_path],
+                        entrypoint_module="main",
+                        entrypoint_object="app",
+                        requirements_file=requirements_file_path,
+                        class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+                    ),
+                )
+                request_mock.assert_called_with(
+                    "post",
+                    "reasoningEngines",
+                    {
+                        "displayName": _TEST_AGENT_ENGINE_DISPLAY_NAME,
+                        "description": _TEST_AGENT_ENGINE_DESCRIPTION,
+                        "spec": {
+                            "agent_framework": "custom",
+                            "source_code_spec": {
+                                "inline_source": {"source_archive": "test_tarball"},
+                                "python_spec": {
+                                    "version": _TEST_PYTHON_VERSION,
+                                    "entrypoint_module": "main",
+                                    "entrypoint_object": "app",
+                                    "requirements_file": requirements_file_path,
+                                },
+                            },
+                            "class_methods": _TEST_AGENT_ENGINE_CLASS_METHODS,
+                        },
+                    },
+                    None,
+                )
+                mock_create_base64_encoded_tarball.assert_called_once_with(
+                    source_packages=[test_file_path]
+                )
+
     @mock.patch.object(agent_engines.AgentEngines, "_create_config")
     @mock.patch.object(_agent_engines_utils, "_await_operation")
     def test_create_agent_engine_with_class_methods(
@@ -1613,6 +1775,10 @@ def test_create_agent_engine_with_class_methods(
                 labels=None,
                 agent_server_mode=None,
                 class_methods=_TEST_AGENT_ENGINE_CLASS_METHODS,
+                source_packages=None,
+                entrypoint_module=None,
+                entrypoint_object=None,
+                requirements_file=None,
             )
             request_mock.assert_called_with(
                 "post",
@@ -1772,9 +1938,9 @@ def test_update_agent_engine_env_vars(
                 [
                     "spec.package_spec.pickle_object_gcs_uri",
                     "spec.package_spec.requirements_gcs_uri",
+                    "spec.class_methods",
                     "spec.deployment_spec.env",
                     "spec.deployment_spec.secret_env",
-                    "spec.class_methods",
                     "spec.agent_framework",
                 ]
             )

vertexai/_genai/_agent_engines_utils.py

@@ -16,6 +16,7 @@
 
 import abc
 import asyncio
+import base64
 from importlib import metadata as importlib_metadata
 import inspect
 import io
@@ -1161,6 +1162,30 @@ def _upload_extra_packages(
     logger.info(f"Writing to {dir_name}/{_EXTRA_PACKAGES_FILE}")
 
 
+def _create_base64_encoded_tarball(
+    *,
+    source_packages: Sequence[str],
+) -> str:
+    """Creates a base64 encoded tarball from the source packages."""
+    logger.info("Creating in-memory tarfile of source_packages")
+    tar_fileobj = io.BytesIO()
+    project_dir = os.path.realpath(os.getcwd())
+    with tarfile.open(fileobj=tar_fileobj, mode="w|gz") as tar:
+        for file in source_packages:
+            real_file_path = os.path.realpath(file)
+            if real_file_path != project_dir and not real_file_path.startswith(
+                project_dir + os.sep
+            ):
+                raise ValueError(
+                    f"File path '{file}' is outside the project directory "
+                    f"'{project_dir}'."
+                )
+            tar.add(file)
+    tar_fileobj.seek(0)
+    tarball_bytes = tar_fileobj.read()
+    return base64.b64encode(tarball_bytes).decode("utf-8")
+
+
 def _validate_extra_packages_or_raise(
     *,
     extra_packages: Sequence[str],