diff --git a/skyvern/services/script_service.py b/skyvern/services/script_service.py index c98b5e13..71c664bb 100644 --- a/skyvern/services/script_service.py +++ b/skyvern/services/script_service.py @@ -27,6 +27,7 @@ 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.schemas.workflow_runs import WorkflowRunBlock from skyvern.forge.sdk.workflow.models.block import ( CodeBlock, FileParserBlock, @@ -271,6 +272,33 @@ async def execute_script( LOG.info("Script executed successfully", script_id=script_id) +async def _take_workflow_run_block_screenshot( + workflow_run_id: str, + organization_id: str, + workflow_run_block: WorkflowRunBlock, +) -> None: + """ + This function is a copy of the block screenshot logic from the execute_safe function in the block.py file. + """ + browser_state = app.BROWSER_MANAGER.get_for_workflow_run(workflow_run_id) + if not browser_state: + LOG.warning("No browser state found when creating workflow_run_block", workflow_run_id=workflow_run_id) + else: + screenshot = await browser_state.take_fullpage_screenshot( + use_playwright_fullpage=app.EXPERIMENTATION_PROVIDER.is_feature_enabled_cached( + "ENABLE_PLAYWRIGHT_FULLPAGE", + workflow_run_id, + properties={"organization_id": str(organization_id)}, + ) + ) + if screenshot: + await app.ARTIFACT_MANAGER.create_workflow_run_block_artifact( + workflow_run_block=workflow_run_block, + artifact_type=ArtifactType.SCREENSHOT_LLM, + data=screenshot, + ) + + async def _create_workflow_block_run_and_task( block_type: BlockType, prompt: str | None = None, @@ -343,6 +371,12 @@ async def _create_workflow_block_run_and_task( organization_id=organization_id, ) + await _take_workflow_run_block_screenshot( + workflow_run_id=workflow_run_id, + organization_id=organization_id, + workflow_run_block=workflow_run_block, + ) + context.step_id = step_id context.task_id = task_id