Skip to content

Commit

Permalink
keep downloaded files after emailed (Skyvern-AI#1322)
Browse files Browse the repository at this point in the history
  • Loading branch information
LawyZheng authored Dec 5, 2024
1 parent fc56299 commit 0099e78
Showing 1 changed file with 48 additions and 53 deletions.
101 changes: 48 additions & 53 deletions skyvern/forge/sdk/workflow/models/block.py
Original file line number Diff line number Diff line change
Expand Up @@ -1271,66 +1271,61 @@ async def _build_email_message(
file_names_by_hash: dict[str, list[str]] = defaultdict(list)

for filename in self._get_file_paths(workflow_run_context, workflow_run_id):
path = None
try:
if filename.startswith("s3://"):
path = await download_from_s3(self.get_async_aws_client(), filename)
elif filename.startswith("http://") or filename.startswith("https://"):
path = await download_file(filename)
else:
LOG.info("SendEmailBlock: Looking for file locally", filename=filename)
if not os.path.exists(filename):
raise FileNotFoundError(f"File not found: {filename}")
if not os.path.isfile(filename):
raise IsADirectoryError(f"Path is a directory: {filename}")

LOG.info("SendEmailBlock: Found file locally", path=path)
path = filename

if not path:
if filename.startswith("s3://"):
path = await download_from_s3(self.get_async_aws_client(), filename)
elif filename.startswith("http://") or filename.startswith("https://"):
path = await download_file(filename)
else:
LOG.info("SendEmailBlock: Looking for file locally", filename=filename)
if not os.path.exists(filename):
raise FileNotFoundError(f"File not found: {filename}")
if not os.path.isfile(filename):
raise IsADirectoryError(f"Path is a directory: {filename}")

path = filename
LOG.info("SendEmailBlock: Found file locally", path=path)

if not path:
raise FileNotFoundError(f"File not found: {filename}")

# Guess the content type based on the file's extension. Encoding
# will be ignored, although we should check for simple things like
# gzip'd or compressed files.
kind = filetype.guess(path)
if kind:
ctype = kind.mime
extension = kind.extension
else:
# No guess could be made, or the file is encoded (compressed), so
# use a generic bag-of-bits type.
ctype = "application/octet-stream"
extension = None

# Guess the content type based on the file's extension. Encoding
# will be ignored, although we should check for simple things like
# gzip'd or compressed files.
kind = filetype.guess(path)
if kind:
ctype = kind.mime
extension = kind.extension
else:
# No guess could be made, or the file is encoded (compressed), so
# use a generic bag-of-bits type.
ctype = "application/octet-stream"
extension = None

maintype, subtype = ctype.split("/", 1)
attachment_path = Path(path)
attachment_filename = attachment_path.name
maintype, subtype = ctype.split("/", 1)
attachment_path = Path(path)
attachment_filename = attachment_path.name

# Check if the filename has an extension
if not attachment_path.suffix:
# If no extension, guess it based on the MIME type
if extension:
attachment_filename += f".{extension}"
# Check if the filename has an extension
if not attachment_path.suffix:
# If no extension, guess it based on the MIME type
if extension:
attachment_filename += f".{extension}"

LOG.info(
"SendEmailBlock: Adding attachment",
filename=attachment_filename,
LOG.info(
"SendEmailBlock: Adding attachment",
filename=attachment_filename,
maintype=maintype,
subtype=subtype,
)
with open(path, "rb") as fp:
msg.add_attachment(
fp.read(),
maintype=maintype,
subtype=subtype,
filename=attachment_filename,
)
with open(path, "rb") as fp:
msg.add_attachment(
fp.read(),
maintype=maintype,
subtype=subtype,
filename=attachment_filename,
)
file_hash = calculate_sha256_for_file(path)
file_names_by_hash[file_hash].append(path)
finally:
if path:
os.unlink(path)
file_hash = calculate_sha256_for_file(path)
file_names_by_hash[file_hash].append(path)

# Calculate file stats based on content hashes
total_files = sum(len(files) for files in file_names_by_hash.values())
Expand Down

0 comments on commit 0099e78

Please sign in to comment.