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

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

@@ -15,10 +15,10 @@
 
 # Code generated by the Google Gen AI SDK generator DO NOT EDIT.
 
-import base64
 import functools
 import json
 import logging
+import mimetypes
 from typing import Any, Iterator, Optional, Union
 from urllib.parse import urlencode
 
@@ -618,20 +618,52 @@ 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", "")
+            input_chunks.append(
+                types.Chunk(
+                    mime_type=file.get("mimeType", ""),
+                    data=file.get("content", b""),
+                    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 is None:
+                # if mime_type is not available, try to guess the mime_type from the file_name.
+                if (
+                    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
+
+            output_chunks.append(output)
+
+        response = types.ExecuteSandboxEnvironmentResponse(outputs=output_chunks)
+
         return response
 
     def get(