Skip to content

Commit ddc6982

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
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
1 parent 0d1240e commit ddc6982

File tree

2 files changed

+38
-10
lines changed

2 files changed

+38
-10
lines changed

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

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,15 +33,30 @@ def test_execute_code_sandbox(client):
3333
config=types.CreateAgentEngineSandboxConfig(display_name="test_sandbox"),
3434
)
3535
assert isinstance(operation, types.AgentEngineSandboxOperation)
36+
37+
code = """
38+
with open("test.txt", "r") as input:
39+
with open("output.txt", "w") as output_txt:
40+
for line in input:
41+
output_txt.write(line)
42+
"""
3643
input_data = {
37-
"language": "python",
38-
"code": 'with open("hello.txt","w") as file:\n file.write("Hello, world!")',
44+
"code": code,
45+
"files": [
46+
{
47+
"name": "test.txt",
48+
"mimeType": "text/plain",
49+
"content": b"Hello, world!",
50+
}
51+
],
3952
}
4053
response = client.agent_engines.sandboxes.execute_code(
4154
name=operation.response.name,
4255
input_data=input_data,
4356
)
4457
assert response.outputs[0].mime_type == "application/json"
58+
assert response.outputs[1].data == b"Hello, world!"
59+
assert response.outputs[1].metadata.attributes.get("file_name") == b"output.txt"
4560

4661

4762
pytestmark = pytest_helper.setup(

vertexai/_genai/sandboxes.py

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
# Code generated by the Google Gen AI SDK generator DO NOT EDIT.
1717

18-
import base64
1918
import functools
2019
import json
2120
import logging
@@ -618,17 +617,31 @@ def execute_code(
618617
Returns:
619618
ExecuteSandboxEnvironmentResponse: The response from executing the code.
620619
"""
621-
json_string = json.dumps(input_data)
622-
623-
base64_bytes = base64.b64encode(json_string.encode("utf-8"))
624-
base64_string = base64_bytes.decode("utf-8")
620+
input_chunks = []
621+
622+
if input_data.get("code") is not None:
623+
code = input_data.get("code", "")
624+
json_code = json.dumps({"code": code}).encode("utf-8")
625+
input_chunks.append(
626+
types.Chunk(
627+
mime_type="application/json",
628+
data=json_code,
629+
)
630+
)
625631

626-
# Only single JSON input is supported for now.
627-
inputs = [{"mime_type": "application/json", "data": base64_string}]
632+
for file in input_data.get("files", []):
633+
file_name = file.get("name", "")
634+
input_chunks.append(
635+
types.Chunk(
636+
mime_type=file.get("mimeType", ""),
637+
data=file.get("content", b""),
638+
metadata={"attributes": {"file_name": file_name.encode("utf-8")}},
639+
)
640+
)
628641

629642
response = self._execute_code(
630643
name=name,
631-
inputs=inputs,
644+
inputs=input_chunks,
632645
config=config,
633646
)
634647

0 commit comments

Comments
 (0)