add new workflow block (#1228)

This commit is contained in:
LawyZheng
2024-11-21 15:12:26 +08:00
committed by GitHub
parent 4271ca9ecf
commit 3f209404f7
16 changed files with 483 additions and 13 deletions

View File

@@ -10,7 +10,7 @@ from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from skyvern.config import settings
from skyvern.exceptions import WorkflowParameterNotFound
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType, TaskPromptTemplate
from skyvern.forge.sdk.db.exceptions import NotFoundError
from skyvern.forge.sdk.db.models import (
ActionModel,
@@ -97,6 +97,8 @@ class AgentDB:
self,
url: str,
title: str | None,
complete_criterion: str | None,
terminate_criterion: str | None,
navigation_goal: str | None,
data_extraction_goal: str | None,
navigation_payload: dict[str, Any] | list | str | None,
@@ -111,17 +113,21 @@ class AgentDB:
retry: int | None = None,
max_steps_per_run: int | None = None,
error_code_mapping: dict[str, str] | None = None,
prompt_template: str = TaskPromptTemplate.ExtractAction,
) -> Task:
try:
async with self.Session() as session:
new_task = TaskModel(
status="created",
prompt_template=prompt_template,
url=url,
title=title,
webhook_callback_url=webhook_callback_url,
totp_verification_url=totp_verification_url,
totp_identifier=totp_identifier,
navigation_goal=navigation_goal,
complete_criterion=complete_criterion,
terminate_criterion=terminate_criterion,
data_extraction_goal=data_extraction_goal,
navigation_payload=navigation_payload,
organization_id=organization_id,

View File

@@ -3,3 +3,19 @@ from enum import StrEnum
class OrganizationAuthTokenType(StrEnum):
api = "api"
class TaskPromptTemplate(StrEnum):
ExtractAction = "extract-action"
DecisiveCriterionValidate = "decisive-criterion-validate"
SingleClickAction = "single-click-action"
SingleInputAction = "single-input-action"
SingleUploadAction = "single-upload-action"
SingleSelectAction = "single-select-action"
class ActionType(StrEnum):
Click = "CLICK"
InputText = "INPUT_TEXT"
UploadFile = "UPLOAD_FILE"
SelectOption = "SELECT_OPTION"

View File

@@ -17,7 +17,7 @@ from sqlalchemy import (
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import DeclarativeBase
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType
from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType, TaskPromptTemplate
from skyvern.forge.sdk.db.id import (
generate_action_id,
generate_artifact_id,
@@ -54,9 +54,12 @@ class TaskModel(Base):
totp_verification_url = Column(String)
totp_identifier = Column(String)
title = Column(String)
prompt_template = Column(String, default=TaskPromptTemplate.ExtractAction)
url = Column(String)
navigation_goal = Column(String)
data_extraction_goal = Column(String)
complete_criterion = Column(String)
terminate_criterion = Column(String)
navigation_payload = Column(JSON)
extracted_information = Column(JSON)
failure_reason = Column(String)

View File

@@ -60,8 +60,11 @@ def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False) -> Task:
status=TaskStatus(task_obj.status),
created_at=task_obj.created_at,
modified_at=task_obj.modified_at,
prompt_template=task_obj.prompt_template,
title=task_obj.title,
url=task_obj.url,
complete_criterion=task_obj.complete_criterion,
terminate_criterion=task_obj.terminate_criterion,
webhook_callback_url=task_obj.webhook_callback_url,
totp_verification_url=task_obj.totp_verification_url,
totp_identifier=task_obj.totp_identifier,

View File

@@ -104,6 +104,14 @@ class Step(BaseModel):
return False
def is_success(self) -> bool:
if self.status != StepStatus.completed:
return False
# TODO (kerem): Remove this check once we have backfilled all the steps
if self.output is None or self.output.actions_and_results is None:
return False
return True
def is_terminated(self) -> bool:
if self.status != StepStatus.completed:
return False

View File

@@ -8,6 +8,7 @@ from pydantic import BaseModel, Field, HttpUrl, field_validator
from skyvern.exceptions import BlockedHost, InvalidTaskStatusTransition, TaskAlreadyCanceled
from skyvern.forge.sdk.core.validators import is_blocked_host
from skyvern.forge.sdk.db.enums import TaskPromptTemplate
class ProxyLocation(StrEnum):
@@ -77,6 +78,19 @@ class TaskBase(BaseModel):
default=None,
description="The requested schema of the extracted information.",
)
complete_criterion: str | None = Field(
default=None, description="Criterion to complete", examples=["Complete if 'hello world' shows up on the page"]
)
terminate_criterion: str | None = Field(
default=None,
description="Criterion to terminate",
examples=["Terminate if 'existing account' shows up on the page"],
)
prompt_template: str | None = Field(
default=TaskPromptTemplate.ExtractAction,
description="The prompt template used for task",
examples=[TaskPromptTemplate.ExtractAction, TaskPromptTemplate.DecisiveCriterionValidate],
)
class TaskRequest(TaskBase):

View File

@@ -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")]

View File

@@ -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")]

View File

@@ -18,6 +18,7 @@ from skyvern.forge.sdk.artifact.models import ArtifactType
from skyvern.forge.sdk.core import skyvern_context
from skyvern.forge.sdk.core.security import generate_skyvern_signature
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
from skyvern.forge.sdk.db.enums import ActionType, TaskPromptTemplate
from skyvern.forge.sdk.models import Organization, Step
from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task
from skyvern.forge.sdk.workflow.exceptions import (
@@ -28,6 +29,7 @@ from skyvern.forge.sdk.workflow.exceptions import (
WorkflowParameterMissingRequiredValue,
)
from skyvern.forge.sdk.workflow.models.block import (
ActionBlock,
BlockStatus,
BlockType,
BlockTypeVar,
@@ -39,6 +41,7 @@ from skyvern.forge.sdk.workflow.models.block import (
TaskBlock,
TextPromptBlock,
UploadToS3Block,
ValidationBlock,
)
from skyvern.forge.sdk.workflow.models.parameter import (
PARAMETER_TYPE,
@@ -1333,4 +1336,65 @@ class WorkflowService:
file_type=block_yaml.file_type,
continue_on_failure=block_yaml.continue_on_failure,
)
elif block_yaml.block_type == BlockType.VALIDATION:
validation_block_parameters = (
[parameters[parameter_key] for parameter_key in block_yaml.parameter_keys]
if block_yaml.parameter_keys
else []
)
if not block_yaml.complete_criterion and not block_yaml.terminate_criterion:
raise Exception("Both complete criterion and terminate criterion are empty")
return ValidationBlock(
label=block_yaml.label,
prompt_template=TaskPromptTemplate.DecisiveCriterionValidate,
parameters=validation_block_parameters,
output_parameter=output_parameter,
complete_criterion=block_yaml.complete_criterion,
terminate_criterion=block_yaml.terminate_criterion,
error_code_mapping=block_yaml.error_code_mapping,
continue_on_failure=block_yaml.continue_on_failure,
# only need one step for validation block
max_steps_per_run=1,
)
elif block_yaml.block_type == BlockType.ACTION:
action_block_parameters = (
[parameters[parameter_key] for parameter_key in block_yaml.parameter_keys]
if block_yaml.parameter_keys
else []
)
prompt_template = ""
if block_yaml.action_type == ActionType.Click:
prompt_template = TaskPromptTemplate.SingleClickAction
elif block_yaml.action_type == ActionType.InputText:
prompt_template = TaskPromptTemplate.SingleInputAction
elif block_yaml.action_type == ActionType.UploadFile:
prompt_template = TaskPromptTemplate.SingleUploadAction
elif block_yaml.action_type == ActionType.SelectOption:
prompt_template = TaskPromptTemplate.SingleSelectAction
if not prompt_template:
raise Exception("not supported action type for action block")
return ActionBlock(
prompt_template=prompt_template,
label=block_yaml.label,
url=block_yaml.url,
title=block_yaml.title,
parameters=action_block_parameters,
output_parameter=output_parameter,
navigation_goal=block_yaml.navigation_goal,
error_code_mapping=block_yaml.error_code_mapping,
max_retries=block_yaml.max_retries,
complete_on_download=block_yaml.complete_on_download,
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,
max_steps_per_run=1,
)
raise ValueError(f"Invalid block type {block_yaml.block_type}")