The saved filename should include the original file name instead of a generated one.

PiperOrigin-RevId: 761727488
This commit is contained in:
Google Team Member 2025-05-21 17:09:02 -07:00 committed by Copybara-Service
parent d7af2b2815
commit 505d936007

View File

@ -148,18 +148,15 @@ class VertexAiCodeExecutor(BaseCodeExecutor):
)
# Save output file as artifacts.
current_timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
file_name_prefix = '%s_' % str(current_timestamp)
saved_files = []
file_count = 0
for output_file in code_execution_result['output_files']:
file_type = output_file['name'].split('.')[-1]
file_name = file_name_prefix + '%d.%s' % (file_count, file_type)
if file_type in _SUPPORTED_IMAGE_TYPES:
file_count += 1
saved_files.append(
File(
name='plot_' + file_name,
name=output_file['name'],
content=output_file['contents'],
mime_type=f'image/{file_type}',
)
@ -168,16 +165,16 @@ class VertexAiCodeExecutor(BaseCodeExecutor):
file_count += 1
saved_files.append(
File(
name='data_' + file_name,
name=output_file['name'],
content=output_file['contents'],
mime_type=f'text/{file_type}',
)
)
else:
mime_type, _ = mimetypes.guess_type(file_name)
mime_type, _ = mimetypes.guess_type(output_file['name'])
saved_files.append(
File(
name=file_name,
name=output_file['name'],
content=output_file['contents'],
mime_type=mime_type,
)