add new workflow block (#1228)

skyvern/forge/sdk/workflow/models/block.py

@@ -40,6 +40,7 @@ from skyvern.forge.sdk.api.files import (
     get_path_for_workflow_download_directory,
 )
 from skyvern.forge.sdk.api.llm.api_handler_factory import LLMAPIHandlerFactory
+from skyvern.forge.sdk.db.enums import TaskPromptTemplate
 from skyvern.forge.sdk.schemas.tasks import Task, TaskOutput, TaskStatus
 from skyvern.forge.sdk.settings_manager import SettingsManager
 from skyvern.forge.sdk.workflow.context_manager import WorkflowRunContext
@@ -69,6 +70,8 @@ class BlockType(StrEnum):
     UPLOAD_TO_S3 = "upload_to_s3"
     SEND_EMAIL = "send_email"
     FILE_URL_PARSER = "file_url_parser"
+    VALIDATION = "validation"
+    ACTION = "action"
 
 
 class BlockStatus(StrEnum):
@@ -174,11 +177,12 @@ class Block(BaseModel, abc.ABC):
         pass
 
 
-class TaskBlock(Block):
-    block_type: Literal[BlockType.TASK] = BlockType.TASK
-
+class BaseTaskBlock(Block):
+    prompt_template: str = TaskPromptTemplate.ExtractAction
     url: str | None = None
     title: str = ""
+    complete_criterion: str | None = None
+    terminate_criterion: str | None = None
     navigation_goal: str | None = None
     data_extraction_goal: str | None = None
     data_schema: dict[str, Any] | list | None = None
@@ -464,6 +468,10 @@ class TaskBlock(Block):
         )
 
 
+class TaskBlock(BaseTaskBlock):
+    block_type: Literal[BlockType.TASK] = BlockType.TASK
+
+
 class ForLoopBlock(Block):
     block_type: Literal[BlockType.FOR_LOOP] = BlockType.FOR_LOOP
 
@@ -1264,6 +1272,36 @@ class FileParserBlock(Block):
         )
 
 
+class ValidationBlock(BaseTaskBlock):
+    block_type: Literal[BlockType.VALIDATION] = BlockType.VALIDATION
+
+    def get_all_parameters(
+        self,
+        workflow_run_id: str,
+    ) -> list[PARAMETER_TYPE]:
+        return self.parameters
+
+    async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult:
+        task_order, _ = await self.get_task_order(workflow_run_id, 0)
+        is_first_task = task_order == 0
+        if is_first_task:
+            return self.build_block_result(
+                success=False,
+                failure_reason="Validation block should not be the first block",
+                output_parameter_value=None,
+                status=BlockStatus.terminated,
+            )
+
+        return await super().execute(workflow_run_id=workflow_run_id, kwargs=kwargs)
+
+
+class ActionBlock(BaseTaskBlock):
+    block_type: Literal[BlockType.ACTION] = BlockType.ACTION
+
+    async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult:
+        return await super().execute(workflow_run_id=workflow_run_id, kwargs=kwargs)
+
+
 BlockSubclasses = Union[
     ForLoopBlock,
     TaskBlock,
@@ -1273,5 +1311,7 @@ BlockSubclasses = Union[
     UploadToS3Block,
     SendEmailBlock,
     FileParserBlock,
+    ValidationBlock,
+    ActionBlock,
 ]
 BlockTypeVar = Annotated[BlockSubclasses, Field(discriminator="block_type")]

skyvern/forge/sdk/workflow/models/yaml.py

@@ -3,6 +3,7 @@ from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Field
 
+from skyvern.forge.sdk.db.enums import ActionType
 from skyvern.forge.sdk.schemas.tasks import ProxyLocation
 from skyvern.forge.sdk.workflow.models.block import BlockType, FileType
 from skyvern.forge.sdk.workflow.models.parameter import ParameterType, WorkflowParameterType
@@ -208,6 +209,32 @@ class FileParserBlockYAML(BlockYAML):
     file_type: FileType
 
 
+class ValidationBlockYAML(BlockYAML):
+    block_type: Literal[BlockType.VALIDATION] = BlockType.VALIDATION  # type: ignore
+
+    complete_criterion: str | None = None
+    terminate_criterion: str | None = None
+    error_code_mapping: dict[str, str] | None = None
+    parameter_keys: list[str] | None = None
+
+
+class ActionBlockYAML(BlockYAML):
+    action_type: ActionType
+    block_type: Literal[BlockType.ACTION] = BlockType.ACTION  # type: ignore
+
+    url: str | None = None
+    title: str = ""
+    navigation_goal: str | None = None
+    error_code_mapping: dict[str, str] | None = None
+    max_retries: int = 0
+    parameter_keys: list[str] | None = None
+    complete_on_download: bool = False
+    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
@@ -228,6 +255,8 @@ BLOCK_YAML_SUBCLASSES = (
     | UploadToS3BlockYAML
     | SendEmailBlockYAML
     | FileParserBlockYAML
+    | ValidationBlockYAML
+    | ActionBlockYAML
 )
 BLOCK_YAML_TYPES = Annotated[BLOCK_YAML_SUBCLASSES, Field(discriminator="block_type")]