feat: GenAI SDK client - Support agent engine sandbox http request in genai sdk

PiperOrigin-RevId: 816865842

Author: vertex-sdk-bot

.kokoro/build.sh

@@ -37,7 +37,7 @@ export PROJECT_ID=$(cat "${KOKORO_GFILE_DIR}/project-id.json")
 python3 -m pip uninstall --yes --quiet nox-automation
 
 # Install nox
-python3 -m pip install --upgrade --quiet nox
+python3 -m pip install --upgrade --quiet uv nox
 python3 -m nox --version
 
 # If this is a continuous build, send the test log to the FlakyBot.

.kokoro/requirements.in

@@ -6,6 +6,7 @@ twine
 wheel
 setuptools
 nox>=2022.11.21 # required to remove dependency on py
+uv
 charset-normalizer<3
 click<8.1.0
 build

.kokoro/requirements.txt

@@ -1,5 +1,5 @@
 #
-# This file is autogenerated by pip-compile with Python 3.9
+# This file is autogenerated by pip-compile with Python 3.13
 # by the following command:
 #
 #    pip-compile --allow-unsafe --generate-hashes requirements.in
@@ -255,8 +255,6 @@ importlib-metadata==6.8.0 \
     --hash=sha256:dbace7892d8c0c4ac1ad096662232f831d4e64f4c4545bd53016a3e9d4654743
     # via
     #   -r requirements.in
-    #   build
-    #   keyring
     #   twine
 jaraco-classes==3.3.0 \
     --hash=sha256:10afa92b6743f25c0cf5f37c6bb6e18e2c5bb84a16527ccfc0040ea377e7aaeb \
@@ -491,12 +489,6 @@ six==1.16.0 \
     # via
     #   gcp-docuploader
     #   python-dateutil
-tomli==2.0.1 \
-    --hash=sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc \
-    --hash=sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f
-    # via
-    #   build
-    #   pyproject-hooks
 twine==4.0.2 \
     --hash=sha256:929bc3c280033347a00f847236564d1c52a3e61b1ac2516c97c48f3ceab756d8 \
     --hash=sha256:9e102ef5fdd5a20661eb88fad46338806c3bd32cf1db729603fe3697b1bc83c8
@@ -511,6 +503,27 @@ urllib3==2.0.7 \
     # via
     #   requests
     #   twine
+uv==0.9.5 \
+    --hash=sha256:02db727beb94a2137508cee5a785c3465d150954ca9abdff2d8157c76dea163e \
+    --hash=sha256:0316493044035098666d6e99c14bd61b352555d9717d57269f4ce531855330fa \
+    --hash=sha256:133e2614e1ff3b34c2606595d8ae55710473ebb7516bfa5708afc00315730cd1 \
+    --hash=sha256:21452ece590ddb90e869a478ca4c2ba70be180ec0d6716985ee727b9394c8aa5 \
+    --hash=sha256:3a4ecbfdcbd3dae4190428874762c791e05d2c97ff2872bf6c0a30ed5c4ea9ca \
+    --hash=sha256:3eb31c9896dc2c88f6a9f1d693be2409fe2fc2e3d90827956e4341c2b2171289 \
+    --hash=sha256:48a12390421f91af8a8993cf15c38297c0bb121936046286e287975b2fbf1789 \
+    --hash=sha256:48a3542835d37882ff57d1ff91b757085525d98756712fa61cf9941d3dda8ebf \
+    --hash=sha256:5bb4996329ba47e7e775baba4a47e85092aa491d708a66e63b564e9b306bfb7e \
+    --hash=sha256:6452eb6257e37e1ebd97430b5f5e10419da2c3ca35b4086540ec4163b4b2f25c \
+    --hash=sha256:6507bbbcd788553ec4ad5a96fa19364dc0f58b023e31d79868773559a83ec181 \
+    --hash=sha256:6a046c2e833169bf26f461286aab58a2ba8d48ed2220bfcf119dcfaf87163116 \
+    --hash=sha256:8603bb902e578463c50c3ddd4ee376ba4172ccdf4979787f8948747d1bb0e18b \
+    --hash=sha256:922cd784cce36bbdc7754b590d28c276698c85791c18cd4c6a7e917db4480440 \
+    --hash=sha256:9fc13b4b943d19adac52d7dcd2159e96ab2e837ac49a79e20714ed25f1f1b7f9 \
+    --hash=sha256:c465f2e342cab908849b8ce83e14fd4cf75f5bed55802d0acf1399f9d02f92d9 \
+    --hash=sha256:c966e3a4fe4de3b0a6279d0a835c79f9cddbb3693f52d140910cbbed177c5742 \
+    --hash=sha256:d8835d2c034421ac2235fb658bb4f669a301a0f1eb00a8430148dd8461b65641 \
+    --hash=sha256:f8eb34ebebac4b45334ce7082cca99293b71fb32b164651f1727c8a640e5b387
+    # via -r requirements.in
 virtualenv==20.24.6 \
     --hash=sha256:02ece4f56fbf939dbbc33c0715159951d6bf14aaf5457b092e4548e1382455af \
     --hash=sha256:520d056652454c5098a00c0f073611ccbea4c79089331f60bf9d7ba247bb7381

.kokoro/test-samples-impl.sh

@@ -34,7 +34,7 @@ env | grep KOKORO
 
 # Install nox
 # `virtualenv==20.26.6` is added for Python 3.7 compatibility
-python3 -m pip install --upgrade --quiet nox virtualenv==20.26.6
+python3 -m pip install --upgrade --quiet nox uv virtualenv==20.26.6
 
 # Use secrets acessor service account to get secrets
 if [[ -f "${KOKORO_GFILE_DIR}/secrets_viewer_service_account.json" ]]; then

noxfile.py

@@ -25,6 +25,8 @@
 
 import nox
 
+nox.options.default_venv_backend = "uv"
+
 FLAKE8_VERSION = "flake8==6.1.0"
 BLACK_VERSION = "black==24.8.0"
 ISORT_VERSION = "isort==5.10.1"
@@ -161,7 +163,7 @@ def format(session):
 @nox.session(python=DEFAULT_PYTHON_VERSION)
 def lint_setup_py(session):
     """Verify that setup.py is valid (including RST check)."""
-    session.install("docutils", "pygments")
+    session.install("docutils", "pygments", "setuptools")
     session.run("python", "setup.py", "check", "--restructuredtext", "--strict")
 
 

tests/unit/aiplatform/test_training_jobs.py

@@ -662,15 +662,6 @@ def test_packager_creates_and_copies_python_package(self):
             f"{tsp._ROOT_MODULE}-{tsp._SETUP_PY_VERSION}.tar.gz"
         ).is_file()
 
-    def test_created_package_module_is_installable_and_can_be_run(self):
-        tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH)
-        source_dist_path = tsp.package_and_copy(copy_method=local_copy_method)
-        subprocess.check_output(["pip3", "install", source_dist_path])
-        module_output = subprocess.check_output(
-            [source_utils._get_python_executable(), "-m", tsp.module_name]
-        )
-        assert "hello world" in module_output.decode()
-
     def test_requirements_are_in_package(self):
         tsp = source_utils._TrainingScriptPythonPackager(
             _TEST_LOCAL_SCRIPT_FILE_PATH, requirements=_TEST_REQUIREMENTS

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

@@ -0,0 +1,62 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+
+
+def test_send_command_sandbox(client):
+    # agent_engine = client.agent_engines.create()
+    # assert isinstance(agent_engine, types.AgentEngine)
+    # assert isinstance(agent_engine.api_resource, types.ReasoningEngine)
+    # agent_engine_name = "projects/254005681254/locations/us-central1/reasoningEngines/2112984271655272448"
+    # operation = client.agent_engines.sandboxes.create(
+    #     # name=agent_engine.api_resource.name,
+    #     name=agent_engine_name,
+    #     spec={
+    #         "code_execution_environment": {
+    #             "machineConfig": "MACHINE_CONFIG_VCPU4_RAM4GIB"
+    #         }
+    #     },
+    #     config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"),
+    # )
+    # assert isinstance(operation, types.AgentEngineSandboxOperation)
+
+    client._api_client.project = None
+    client._api_client.location = None
+    client._api_client.api_version = None
+    token = client.agent_engines.sandboxes.generate_access_token(
+        service_account_email="sign-verify-jwt@mariner-proxy.iam.gserviceaccount.com",
+        sandbox_id="tenghuil-manual-sandbox-new",
+    )
+    response = client.agent_engines.sandboxes.send_command(
+        http_method="GET",
+        path="/",
+        query_params={},
+        access_token=token,
+        headers={},
+        request_dict={},
+        sandbox_environment=None,
+    )
+    # assert token == "xxx"
+    assert response.body == "Hello World"
+    # assert response.outputs[0].mime_type == "application/json"
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="agent_engines.sandboxes.send_command",
+)

vertexai/_genai/sandboxes.py

@@ -19,11 +19,16 @@
 import json
 import logging
 import mimetypes
+import random
+import time
 from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
+from google import genai
+from google.cloud import iam_credentials_v1
 from google.genai import _api_module
 from google.genai import _common
+from google.genai import types as genai_types
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
 from google.genai.pagers import Pager
@@ -704,6 +709,52 @@ def delete(
         """
         return self._delete(name=name, config=config)
 
+    def generate_access_token(
+        self,
+        service_account_email: str,
+        sandbox_id: str,
+        port: str = "8080",
+        timeout: int = 3600,
+    ) -> str:
+        """Signs a JWT with a Google Cloud service account."""
+        client = iam_credentials_v1.IAMCredentialsClient()
+        name = f"projects/-/serviceAccounts/{service_account_email}"
+        custom_claims = {"port": port, "sandbox_id": sandbox_id}
+        payload = {
+            "iat": int(time.time()),
+            "exp": int(time.time()) + timeout,
+            "iss": service_account_email,
+            "nonce": random.randint(1, 1000000000),
+            "aud": "vmaas-proxy-api",  # default audience for sandbox proxy
+            **custom_claims,
+        }
+        request = iam_credentials_v1.SignJwtRequest(
+            name=name,
+            payload=json.dumps(payload),
+        )
+        response = client.sign_jwt(request=request)
+        return response.signed_jwt
+
+    def send_command(
+        self,
+        http_method: str,
+        path: str,
+        query_params: Any,
+        access_token: str,
+        headers: dict[str, str],
+        request_dict: dict[str, object],
+        sandbox_environment: Optional[types.SandboxEnvironment] = None,
+    ) -> str | None:
+        """Sends a command to the sandbox."""
+        # TODO(tenghuil): Get connection info from sandbox environment when it's ready.
+        endpoint = "https://test-us-central1.autopush-sandbox.vertexai.goog/" + path
+
+        headers = {"Authorization": f"Bearer {access_token}"}
+        http_options = genai_types.HttpOptions(headers=headers, base_url=endpoint)
+        http_client = genai.Client(vertexai=True, http_options=http_options)
+        response = http_client._api_client.request(http_method, path, request_dict)
+        return response
+
 
 class AsyncSandboxes(_api_module.BaseModule):