From 67c54c4040648da6155f82077c9f12f54d5d9708 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Tue, 9 Sep 2025 23:46:45 -0700 Subject: [PATCH] Script gen - support CloudStorageBlock (#3403) --- .../script_generations/generate_script.py | 61 +++++++++++++ skyvern/services/script_service.py | 86 ++++++++++++++++--- 2 files changed, 133 insertions(+), 14 deletions(-) diff --git a/skyvern/core/script_generations/generate_script.py b/skyvern/core/script_generations/generate_script.py index 06c4203e..98daa369 100644 --- a/skyvern/core/script_generations/generate_script.py +++ b/skyvern/core/script_generations/generate_script.py @@ -32,6 +32,7 @@ from skyvern.core.script_generations.generate_workflow_parameters import ( hydrate_input_text_actions_with_field_names, ) from skyvern.forge import app +from skyvern.schemas.workflows import FileStorageType from skyvern.webeye.actions.action_types import ActionType LOG = structlog.get_logger(__name__) @@ -920,6 +921,64 @@ def _build_code_statement(block: dict[str, Any]) -> cst.SimpleStatementLine: return cst.SimpleStatementLine([cst.Expr(cst.Await(call))]) +def _build_file_upload_statement(block: dict[str, Any]) -> cst.SimpleStatementLine: + """Build a skyvern.upload_file statement.""" + args = [ + cst.Arg( + keyword=cst.Name("label"), + value=_value(block.get("label", "")), + whitespace_after_arg=cst.ParenthesizedWhitespace( + indent=True, + last_line=cst.SimpleWhitespace(INDENT), + ), + ), + cst.Arg( + keyword=cst.Name("parameters"), + value=_value(block.get("parameters", None)), + whitespace_after_arg=cst.ParenthesizedWhitespace( + indent=True, + ), + ), + cst.Arg( + keyword=cst.Name("storage_type"), + value=_value(block.get("storage_type", FileStorageType.S3)), + whitespace_after_arg=cst.ParenthesizedWhitespace( + indent=True, + ), + ), + ] + for key in [ + "s3_bucket", + "aws_access_key_id", + "aws_secret_access_key", + "region_name", + "azure_storage_account_name", + "azure_storage_account_key", + "azure_blob_container_name", + "path", + ]: + if block.get(key) is not None: + args.append( + cst.Arg( + keyword=cst.Name(key), + value=_value(block.get(key, "")), + whitespace_after_arg=cst.ParenthesizedWhitespace( + indent=True, + ), + ) + ) + call = cst.Call( + func=cst.Attribute(value=cst.Name("skyvern"), attr=cst.Name("upload_file")), + args=args, + whitespace_before_args=cst.ParenthesizedWhitespace( + indent=True, + last_line=cst.SimpleWhitespace(INDENT), + ), + ) + + return cst.SimpleStatementLine([cst.Expr(cst.Await(call))]) + + def __build_base_task_statement(block_title: str, block: dict[str, Any]) -> list[cst.Arg]: args = [ cst.Arg( @@ -1031,6 +1090,8 @@ def _build_run_fn(blocks: list[dict[str, Any]], wf_req: dict[str, Any]) -> Funct stmt = _build_goto_statement(block) elif block_type == "code": stmt = _build_code_statement(block) + elif block_type == "file_upload": + stmt = _build_file_upload_statement(block) else: # Default case for unknown block types stmt = cst.SimpleStatementLine([cst.Expr(cst.SimpleString(f"# Unknown block type: {block_type}"))]) diff --git a/skyvern/services/script_service.py b/skyvern/services/script_service.py index 8b4a7a87..530bd71d 100644 --- a/skyvern/services/script_service.py +++ b/skyvern/services/script_service.py @@ -5,6 +5,7 @@ import importlib.util import json import os import uuid +from dataclasses import dataclass from datetime import datetime from typing import Any, cast @@ -26,12 +27,12 @@ from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.models import Step, StepStatus from skyvern.forge.sdk.schemas.files import FileInfo from skyvern.forge.sdk.schemas.tasks import Task, TaskOutput, TaskStatus -from skyvern.forge.sdk.workflow.models.block import CodeBlock, TaskBlock -from skyvern.forge.sdk.workflow.models.parameter import PARAMETER_TYPE +from skyvern.forge.sdk.workflow.models.block import CodeBlock, FileUploadBlock, TaskBlock +from skyvern.forge.sdk.workflow.models.parameter import PARAMETER_TYPE, OutputParameter from skyvern.forge.sdk.workflow.models.workflow import Workflow from skyvern.schemas.runs import RunEngine from skyvern.schemas.scripts import CreateScriptResponse, FileEncoding, FileNode, ScriptFileCreate -from skyvern.schemas.workflows import BlockStatus, BlockType +from skyvern.schemas.workflows import BlockStatus, BlockType, FileStorageType LOG = structlog.get_logger(__name__) jinja_sandbox_env = SandboxedEnvironment() @@ -1376,11 +1377,18 @@ def render_template(template: str, data: dict[str, Any] | None = None) -> str: # Non-task-based blocks -async def run_code( - code: str, - label: str | None = None, - parameters: list[PARAMETER_TYPE] | None = None, -) -> dict[str, Any]: +## Non-task-based block helpers +@dataclass +class BlockValidationOutput: + label: str + output_parameter: OutputParameter + workflow_id: str + workflow_run_id: str + organization_id: str + browser_session_id: str | None = None + + +async def _validate_and_get_output_parameter(label: str | None = None) -> BlockValidationOutput: context = skyvern_context.ensure_context() workflow_id = context.workflow_id workflow_run_id = context.workflow_run_id @@ -1399,16 +1407,66 @@ async def run_code( output_parameter = workflow.get_output_parameter(label) if not output_parameter: raise Exception("Output parameter not found") - - code_block = CodeBlock( - code=code, + return BlockValidationOutput( label=label, - parameters=parameters or [], output_parameter=output_parameter, - ) - block_result = await code_block.execute_safe( + workflow_id=workflow_id, workflow_run_id=workflow_run_id, organization_id=organization_id, browser_session_id=browser_session_id, ) + + +async def run_code( + code: str, + label: str | None = None, + parameters: list[PARAMETER_TYPE] | None = None, +) -> dict[str, Any]: + block_validation_output = await _validate_and_get_output_parameter(label) + code_block = CodeBlock( + code=code, + label=block_validation_output.label, + parameters=parameters or [], + output_parameter=block_validation_output.output_parameter, + ) + block_result = await code_block.execute_safe( + workflow_run_id=block_validation_output.workflow_run_id, + organization_id=block_validation_output.organization_id, + browser_session_id=block_validation_output.browser_session_id, + ) return cast(dict[str, Any], block_result.output_parameter_value) + + +async def upload_file( + label: str | None = None, + parameters: list[PARAMETER_TYPE] | None = None, + storage_type: FileStorageType = FileStorageType.S3, + s3_bucket: str | None = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + region_name: str | None = None, + azure_storage_account_name: str | None = None, + azure_storage_account_key: str | None = None, + azure_blob_container_name: str | None = None, + path: str | None = None, +) -> None: + block_validation_output = await _validate_and_get_output_parameter(label) + file_upload_block = FileUploadBlock( + label=block_validation_output.label, + output_parameter=block_validation_output.output_parameter, + parameters=parameters or [], + storage_type=storage_type, + s3_bucket=s3_bucket, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=region_name, + azure_storage_account_name=azure_storage_account_name, + azure_storage_account_key=azure_storage_account_key, + azure_blob_container_name=azure_blob_container_name, + path=path, + ) + await file_upload_block.execute_safe( + workflow_run_id=block_validation_output.workflow_run_id, + organization_id=block_validation_output.organization_id, + browser_session_id=block_validation_output.browser_session_id, + )