diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 4386758d..cf9f1190 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -77,6 +77,7 @@ class BlockType(StrEnum): EXTRACTION = "extraction" LOGIN = "login" WAIT = "wait" + FILE_DOWNLOAD = "file_download" class BlockStatus(StrEnum): @@ -1348,6 +1349,10 @@ class LoginBlock(BaseTaskBlock): block_type: Literal[BlockType.LOGIN] = BlockType.LOGIN +class FileDownloadBlock(BaseTaskBlock): + block_type: Literal[BlockType.FILE_DOWNLOAD] = BlockType.FILE_DOWNLOAD + + BlockSubclasses = Union[ ForLoopBlock, TaskBlock, @@ -1363,5 +1368,6 @@ BlockSubclasses = Union[ ExtractionBlock, LoginBlock, WaitBlock, + FileDownloadBlock, ] BlockTypeVar = Annotated[BlockSubclasses, Field(discriminator="block_type")] diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py index c9fb647d..c5c048ef 100644 --- a/skyvern/forge/sdk/workflow/models/yaml.py +++ b/skyvern/forge/sdk/workflow/models/yaml.py @@ -283,6 +283,22 @@ class WaitBlockYAML(BlockYAML): wait_sec: int = 0 +class FileDownloadBlockYAML(BlockYAML): + block_type: Literal[BlockType.FILE_DOWNLOAD] = BlockType.FILE_DOWNLOAD # type: ignore + + navigation_goal: str + url: str | None = None + title: str = "" + error_code_mapping: dict[str, str] | None = None + max_retries: int = 0 + max_steps_per_run: int | None = None + parameter_keys: list[str] | None = None + download_suffix: str | None = None + totp_verification_url: str | None = None + totp_identifier: str | None = None + cache_actions: bool = False + + PARAMETER_YAML_SUBCLASSES = ( AWSSecretParameterYAML | BitwardenLoginCredentialParameterYAML @@ -309,6 +325,7 @@ BLOCK_YAML_SUBCLASSES = ( | ExtractionBlockYAML | LoginBlockYAML | WaitBlockYAML + | FileDownloadBlockYAML ) BLOCK_YAML_TYPES = Annotated[BLOCK_YAML_SUBCLASSES, Field(discriminator="block_type")] diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 2282341c..42c98565 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -38,6 +38,7 @@ from skyvern.forge.sdk.workflow.models.block import ( CodeBlock, DownloadToS3Block, ExtractionBlock, + FileDownloadBlock, FileParserBlock, ForLoopBlock, LoginBlock, @@ -1474,4 +1475,28 @@ class WorkflowService: output_parameter=output_parameter, ) + elif block_yaml.block_type == BlockType.FILE_DOWNLOAD: + file_download_block_parameters = ( + [parameters[parameter_key] for parameter_key in block_yaml.parameter_keys] + if block_yaml.parameter_keys + else [] + ) + return FileDownloadBlock( + label=block_yaml.label, + url=block_yaml.url, + title=block_yaml.title, + parameters=file_download_block_parameters, + output_parameter=output_parameter, + navigation_goal=block_yaml.navigation_goal, + error_code_mapping=block_yaml.error_code_mapping, + max_steps_per_run=block_yaml.max_steps_per_run, + max_retries=block_yaml.max_retries, + download_suffix=block_yaml.download_suffix, + continue_on_failure=block_yaml.continue_on_failure, + totp_verification_url=block_yaml.totp_verification_url, + totp_identifier=block_yaml.totp_identifier, + cache_actions=block_yaml.cache_actions, + complete_on_download=True, + ) + raise ValueError(f"Invalid block type {block_yaml.block_type}")