feat: GenAI SDK client - Update input handling inside code execution sandbox

FUTURE_COPYBARA_INTEGRATE_REVIEW=#5883 from googleapis:release-please--branches--main 048acfd
PiperOrigin-RevId: 816441862

Author: vertex-sdk-bot

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

@@ -33,15 +33,30 @@ def test_execute_code_sandbox(client):
         config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"),
     )
     assert isinstance(operation, types.AgentEngineSandboxOperation)
+
+    code = """
+with open("test.txt", "r") as input:
+    with open("output.txt", "w") as output_txt:
+        for line in input:
+            output_txt.write(line)
+"""
     input_data = {
-        "language": "python",
-        "code": 'with open("hello.txt","w") as file:\n file.write("Hello, world!")',
+        "code": code,
+        "files": [
+            {
+                "name": "test.txt",
+                "mimeType": "text/plain",
+                "content": b"Hello, world!",
+            }
+        ],
     }
     response = client.agent_engines.sandboxes.execute_code(
         name=operation.response.name,
         input_data=input_data,
     )
     assert response.outputs[0].mime_type == "application/json"
+    assert response.outputs[1].data == b"Hello, world!"
+    assert response.outputs[1].metadata.attributes.get("file_name") == b"output.txt"
 
 
 pytestmark = pytest_helper.setup(

vertexai/_genai/sandboxes.py

@@ -19,6 +19,7 @@
 import functools
 import json
 import logging
+import mimetypes
 from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
@@ -528,6 +529,22 @@ def _get_sandbox_operation(
         self._api_client._verify_response(return_value)
         return return_value
 
+    _NEEDED_BASE64_ENCODING_MIME_TYPES = [
+        "image/jpeg",
+        "image/png",
+        "image/gif",
+        "image/bmp",
+        "image/webp",
+        "audio/mpeg",
+        "audio/wav",
+        "audio/aac",
+        "audio/ogg",
+        "video/mp4",
+        "video/webm",
+        "video/ogg",
+        "video/mpeg",
+    ]
+
     def create(
         self,
         *,
@@ -618,20 +635,66 @@ def execute_code(
         Returns:
             ExecuteSandboxEnvironmentResponse: The response from executing the code.
         """
-        json_string = json.dumps(input_data)
-
-        base64_bytes = base64.b64encode(json_string.encode("utf-8"))
-        base64_string = base64_bytes.decode("utf-8")
+        input_chunks = []
+
+        if input_data.get("code") is not None:
+            code = input_data.get("code", "")
+            json_code = json.dumps({"code": code}).encode("utf-8")
+            input_chunks.append(
+                types.Chunk(
+                    mime_type="application/json",
+                    data=json_code,
+                )
+            )
 
-        # Only single JSON input is supported for now.
-        inputs = [{"mime_type": "application/json", "data": base64_string}]
+        for file in input_data.get("files", []):
+            file_name = file.get("name", "")
+            mime_type = file.get("mimeType", "")
+            if mime_type is None:
+                mime_type, _ = mimetypes.guess_type(file_name)
+            if mime_type in self._NEEDED_BASE64_ENCODING_MIME_TYPES:
+                base64_bytes = base64.b64encode(file.get("content", b""))
+                content = base64_bytes.decode("utf-8")
+            else:
+                content = file.get("content", b"")
+            input_chunks.append(
+                types.Chunk(
+                    mime_type=mime_type,
+                    data=content,
+                    metadata={"attributes": {"file_name": file_name.encode("utf-8")}},
+                )
+            )
 
         response = self._execute_code(
             name=name,
-            inputs=inputs,
+            inputs=input_chunks,
             config=config,
         )
 
+        output_chunks = []
+        for output in response.outputs:
+            if output.mime_type != "application/json":
+                mime_type = output.mime_type
+                # if mime_type is not available, try to guess the mime_type from the file_name.
+                if (
+                    mime_type is None
+                    and output.metadata is not None
+                    and output.metadata.attributes is not None
+                ):
+                    file_name = output.metadata.attributes.get("file_name", b"").decode(
+                        "utf-8"
+                    )
+                    mime_type, _ = mimetypes.guess_type(file_name)
+                    output.mime_type = mime_type
+
+                # if the mime_type is in the list of mime_types that need base64 encoding,
+                # decode the data.
+                if mime_type in self._NEEDED_BASE64_ENCODING_MIME_TYPES:
+                    output.data = base64.b64decode(output.data)
+            output_chunks.append(output)
+
+        response = types.ExecuteSandboxEnvironmentResponse(outputs=output_chunks)
+
         return response
 
     def get(