mirror of
https://github.com/Skyvern-AI/skyvern.git
synced 2025-09-02 10:41:04 +00:00
125 lines
5.4 KiB
Python
125 lines
5.4 KiB
Python
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
|
|
|
|
|
|
class TextPromptBlockYAML(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.TEXT_PROMPT] = BlockType.TEXT_PROMPT # type: ignore
|
|
|
|
llm_key: str
|
|
prompt: str
|
|
parameter_keys: list[str] | None = None
|
|
json_schema: dict[str, Any] | 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 | TextPromptBlockYAML
|
|
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
|