Give playwright access to CodeBlock (#588)

This commit is contained in:
Kerem Yilmaz
2024-07-11 09:35:07 -07:00
committed by GitHub
parent 48ba239080
commit 87d6e71768
2 changed files with 30 additions and 2 deletions

View File

@@ -1,8 +1,10 @@
import abc import abc
import asyncio
import csv import csv
import json import json
import os import os
import smtplib import smtplib
import textwrap
import uuid import uuid
from dataclasses import dataclass from dataclasses import dataclass
from email.message import EmailMessage from email.message import EmailMessage
@@ -501,8 +503,13 @@ class CodeBlock(Block):
async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult: async def execute(self, workflow_run_id: str, **kwargs: dict) -> BlockResult:
# get workflow run context # get workflow run context
workflow_run_context = self.get_workflow_run_context(workflow_run_id) workflow_run_context = self.get_workflow_run_context(workflow_run_id)
# get all parameters into a dictionary # get all parameters into a dictionary
parameter_values = {} parameter_values = {}
maybe_browser_state = await app.BROWSER_MANAGER.get_for_workflow_run(workflow_run_id)
if maybe_browser_state and maybe_browser_state.page:
parameter_values["skyvern_page"] = maybe_browser_state.page
for parameter in self.parameters: for parameter in self.parameters:
value = workflow_run_context.get_value(parameter.key) value = workflow_run_context.get_value(parameter.key)
secret_value = workflow_run_context.get_original_secret_value_or_none(value) secret_value = workflow_run_context.get_original_secret_value_or_none(value)
@@ -511,9 +518,25 @@ class CodeBlock(Block):
else: else:
parameter_values[parameter.key] = value parameter_values[parameter.key] = value
# Add asyncio and the current event loop to the parameter_values
parameter_values["asyncio"] = asyncio
parameter_values["__builtins__"] = __builtins__ # Include builtins for exec context
local_variables: dict[str, Any] = {} local_variables: dict[str, Any] = {}
exec(self.code, parameter_values, local_variables) result_container: dict[str, Any] = {}
result = {"result": local_variables.get("result")} # Define the user_code function and return it
user_code = textwrap.indent(self.code, " ")
full_code = f"""
async def user_code():
{user_code}
result_container['result'] = locals().get('result')
"""
exec(full_code, {**parameter_values, "result_container": result_container}, local_variables)
# Await the returned user_code function
await local_variables["user_code"]()
result = {"result": result_container.get("result")}
await self.record_output_parameter_value(workflow_run_context, workflow_run_id, result) await self.record_output_parameter_value(workflow_run_context, workflow_run_id, result)
return self.build_block_result(success=True, output_parameter_value=result) return self.build_block_result(success=True, output_parameter_value=result)

View File

@@ -85,6 +85,11 @@ class BrowserManager:
self.pages[workflow_run.workflow_run_id] = browser_state self.pages[workflow_run.workflow_run_id] = browser_state
return browser_state return browser_state
async def get_for_workflow_run(self, workflow_run_id: str) -> BrowserState | None:
if workflow_run_id in self.pages:
return self.pages[workflow_run_id]
return None
def set_video_artifact_for_task(self, task: Task, artifact_id: str) -> None: def set_video_artifact_for_task(self, task: Task, artifact_id: str) -> None:
if task.workflow_run_id and task.workflow_run_id in self.pages: if task.workflow_run_id and task.workflow_run_id in self.pages:
if self.pages[task.workflow_run_id].browser_artifacts.video_artifact_id: if self.pages[task.workflow_run_id].browser_artifacts.video_artifact_id: