Workflow: YAML interface (#123)

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

@@ -0,0 +1,112 @@
+import abc
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field
+
+from skyvern.forge.sdk.workflow.models.block import BlockType
+from skyvern.forge.sdk.workflow.models.parameter import ParameterType, WorkflowParameterType
+
+
+class ParameterYAML(BaseModel, abc.ABC):
+    parameter_type: ParameterType
+    key: str
+    description: str | None = None
+
+
+class AWSSecretParameterYAML(ParameterYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the ParameterType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    parameter_type: Literal[ParameterType.AWS_SECRET] = ParameterType.AWS_SECRET  # type: ignore
+    aws_key: str
+
+
+class WorkflowParameterYAML(ParameterYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the ParameterType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    parameter_type: Literal[ParameterType.WORKFLOW] = ParameterType.WORKFLOW  # type: ignore
+    workflow_parameter_type: WorkflowParameterType
+    default_value: str | int | float | bool | dict | list | None = None
+
+
+class ContextParameterYAML(ParameterYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the ParameterType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    parameter_type: Literal[ParameterType.CONTEXT] = ParameterType.CONTEXT  # type: ignore
+    source_workflow_parameter_key: str
+
+
+class OutputParameterYAML(ParameterYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the ParameterType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    parameter_type: Literal[ParameterType.OUTPUT] = ParameterType.OUTPUT  # type: ignore
+
+
+class BlockYAML(BaseModel, abc.ABC):
+    block_type: BlockType
+    label: str
+    output_parameter_key: str | None = None
+
+
+class TaskBlockYAML(BlockYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the BlockType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    block_type: Literal[BlockType.TASK] = BlockType.TASK  # type: ignore
+
+    url: str | None = None
+    title: str = "Untitled Task"
+    navigation_goal: str | None = None
+    data_extraction_goal: str | None = None
+    data_schema: dict[str, Any] | None = None
+    error_code_mapping: dict[str, str] | None = None
+    max_retries: int = 0
+    parameter_keys: list[str] | None = None
+
+
+class ForLoopBlockYAML(BlockYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the BlockType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    block_type: Literal[BlockType.FOR_LOOP] = BlockType.FOR_LOOP  # type: ignore
+
+    loop_over_parameter_key: str
+    loop_block: BlockYAML
+
+
+class CodeBlockYAML(BlockYAML):
+    # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error:
+    # Parameter 1 of Literal[...] cannot be of type "Any"
+    # This pattern already works in block.py but since the BlockType is not defined in this file, mypy is not able
+    # to infer the type of the parameter_type attribute.
+    block_type: Literal[BlockType.CODE] = BlockType.CODE  # type: ignore
+
+    code: str
+    parameter_keys: list[str] | None = None
+
+
+PARAMETER_YAML_SUBCLASSES = AWSSecretParameterYAML | WorkflowParameterYAML | ContextParameterYAML | OutputParameterYAML
+PARAMETER_YAML_TYPES = Annotated[PARAMETER_YAML_SUBCLASSES, Field(discriminator="parameter_type")]
+
+BLOCK_YAML_SUBCLASSES = TaskBlockYAML | ForLoopBlockYAML | CodeBlockYAML
+BLOCK_YAML_TYPES = Annotated[BLOCK_YAML_SUBCLASSES, Field(discriminator="block_type")]
+
+
+class WorkflowDefinitionYAML(BaseModel):
+    parameters: list[PARAMETER_YAML_TYPES]
+    blocks: list[BLOCK_YAML_TYPES]
+
+
+class WorkflowCreateYAMLRequest(BaseModel):
+    title: str
+    description: str | None = None
+    workflow_definition: WorkflowDefinitionYAML