Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions src/bedrock_agentcore/tools/code_interpreter_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -600,14 +600,15 @@ def install_packages(
def download_file(
self,
path: str,
) -> str:
) -> Union[str, bytes]:
"""Download/read a file from the code interpreter environment.

Args:
path: Path to the file to read.

Returns:
File content as string.
File content as string, or bytes if the file contains binary content
(images, PDFs, etc.).

Raises:
FileNotFoundError: If the file doesn't exist.
Expand All @@ -631,21 +632,26 @@ def download_file(
if "text" in resource:
return resource["text"]
elif "blob" in resource:
return base64.b64decode(resource["blob"]).decode("utf-8")
raw = base64.b64decode(resource["blob"])
try:
return raw.decode("utf-8")
except (UnicodeDecodeError, ValueError):
return raw

raise FileNotFoundError(f"Could not read file: {path}")

def download_files(
self,
paths: List[str],
) -> Dict[str, str]:
) -> Dict[str, Union[str, bytes]]:
"""Download/read multiple files from the code interpreter environment.

Args:
paths: List of file paths to read.

Returns:
Dict mapping file paths to their contents.
Dict mapping file paths to their contents. Values are strings for
text files, or bytes for binary files (images, PDFs, etc.).

Example:
>>> files = client.download_files(['data.csv', 'results.json'])
Expand All @@ -667,7 +673,11 @@ def download_files(
if "text" in resource:
files[file_path] = resource["text"]
elif "blob" in resource:
files[file_path] = base64.b64decode(resource["blob"]).decode("utf-8")
raw = base64.b64decode(resource["blob"])
try:
files[file_path] = raw.decode("utf-8")
except (UnicodeDecodeError, ValueError):
files[file_path] = raw

return files

Expand Down
145 changes: 142 additions & 3 deletions tests/bedrock_agentcore/tools/test_code_interpreter_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1052,8 +1052,44 @@ def test_download_file_binary(self, mock_boto3, mock_get_data_endpoint, mock_get
client.identifier = "test.identifier"
client.session_id = "test-session-id"

original_content = "binary content as text"
encoded_content = base64.b64encode(original_content.encode("utf-8")).decode("utf-8")
binary_content = b"\x89PNG\r\n\x1a\n" # PNG header bytes
encoded_content = base64.b64encode(binary_content).decode("utf-8")

mock_response = {
"stream": [
{
"result": {
"content": [
{"type": "resource", "resource": {"uri": "file://image.png", "blob": encoded_content}}
]
}
}
]
}
client.data_plane_client.invoke_code_interpreter.return_value = mock_response

# Act
result = client.download_file("image.png")

# Assert
assert result == binary_content
assert isinstance(result, bytes)

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
def test_download_file_blob_utf8_returns_str(self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint):
"""Blob content that is valid UTF-8 should be decoded and returned as str."""
# Arrange
mock_session = MagicMock()
mock_session.client.return_value = MagicMock()
mock_boto3.Session.return_value = mock_session
client = CodeInterpreter("us-west-2")
client.identifier = "test.identifier"
client.session_id = "test-session-id"

text_content = "hello world"
encoded_content = base64.b64encode(text_content.encode("utf-8")).decode("utf-8")

mock_response = {
"stream": [
Expand All @@ -1072,7 +1108,8 @@ def test_download_file_binary(self, mock_boto3, mock_get_data_endpoint, mock_get
result = client.download_file("data.bin")

# Assert
assert result == original_content
assert result == text_content
assert isinstance(result, str)

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
Expand Down Expand Up @@ -1152,6 +1189,108 @@ def test_download_files_empty_result(self, mock_boto3, mock_get_data_endpoint, m
# Assert
assert result == {}

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
def test_download_files_binary(self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint):
# Arrange
mock_session = MagicMock()
mock_session.client.return_value = MagicMock()
mock_boto3.Session.return_value = mock_session
client = CodeInterpreter("us-west-2")
client.identifier = "test.identifier"
client.session_id = "test-session-id"

binary_content = b"\x89PNG\r\n\x1a\n" # PNG header bytes
encoded_binary = base64.b64encode(binary_content).decode("utf-8")

mock_response = {
"stream": [
{
"result": {
"content": [
{
"type": "resource",
"resource": {
"uri": "file:///opt/amazon/genesis1p-tools/var/data.csv",
"text": "col1,col2\n1,2",
},
},
{
"type": "resource",
"resource": {
"uri": "file:///opt/amazon/genesis1p-tools/var/chart.png",
"blob": encoded_binary,
},
},
]
}
}
]
}
client.data_plane_client.invoke_code_interpreter.return_value = mock_response

# Act
result = client.download_files(["data.csv", "chart.png"])

# Assert
assert result["/opt/amazon/genesis1p-tools/var/data.csv"] == "col1,col2\n1,2"
assert result["/opt/amazon/genesis1p-tools/var/chart.png"] == binary_content
assert isinstance(result["/opt/amazon/genesis1p-tools/var/chart.png"], bytes)

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
def test_download_files_blob_utf8_returns_str(self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint):
"""Blob content that is valid UTF-8 should be decoded and returned as str in multi-file download."""
# Arrange
mock_session = MagicMock()
mock_session.client.return_value = MagicMock()
mock_boto3.Session.return_value = mock_session
client = CodeInterpreter("us-west-2")
client.identifier = "test.identifier"
client.session_id = "test-session-id"

text_content = "some utf-8 blob content"
encoded_text = base64.b64encode(text_content.encode("utf-8")).decode("utf-8")
binary_content = b"\x89PNG\r\n\x1a\n"
encoded_binary = base64.b64encode(binary_content).decode("utf-8")

mock_response = {
"stream": [
{
"result": {
"content": [
{
"type": "resource",
"resource": {
"uri": "file:///opt/amazon/genesis1p-tools/var/data.bin",
"blob": encoded_text,
},
},
{
"type": "resource",
"resource": {
"uri": "file:///opt/amazon/genesis1p-tools/var/chart.png",
"blob": encoded_binary,
},
},
]
}
}
]
}
client.data_plane_client.invoke_code_interpreter.return_value = mock_response

# Act
result = client.download_files(["data.bin", "chart.png"])

# Assert — UTF-8-valid blob comes back as str, invalid UTF-8 blob comes back as bytes
assert result["/opt/amazon/genesis1p-tools/var/data.bin"] == text_content
assert isinstance(result["/opt/amazon/genesis1p-tools/var/data.bin"], str)
assert result["/opt/amazon/genesis1p-tools/var/chart.png"] == binary_content
assert isinstance(result["/opt/amazon/genesis1p-tools/var/chart.png"], bytes)

@patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
Expand Down
Loading