fix workflow_runs.code_gen to enforce code generate only for task v2 block run (#3597)

This commit is contained in:
Shuchang Zheng
2025-10-02 16:06:54 -07:00
committed by GitHub
parent a7025fc9b7
commit 7e5c4b3a03
14 changed files with 53 additions and 37 deletions

View File

@@ -238,7 +238,6 @@ class Block(BaseModel, abc.ABC):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
pass
@@ -296,7 +295,6 @@ class Block(BaseModel, abc.ABC):
parent_workflow_run_block_id: str | None = None,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
workflow_run_block_id = None
@@ -346,7 +344,6 @@ class Block(BaseModel, abc.ABC):
workflow_run_block_id,
organization_id=organization_id,
browser_session_id=browser_session_id,
code_gen=code_gen,
**kwargs,
)
except Exception as e:
@@ -503,7 +500,6 @@ class BaseTaskBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
workflow_run_context = self.get_workflow_run_context(workflow_run_id)
@@ -1311,7 +1307,6 @@ class ForLoopBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
workflow_run_context = self.get_workflow_run_context(workflow_run_id)
@@ -1489,7 +1484,6 @@ async def wrapper():
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
await app.AGENT_FUNCTION.validate_code_block(organization_id=organization_id)
@@ -1656,7 +1650,6 @@ class TextPromptBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
# Validate block execution
@@ -1741,7 +1734,6 @@ class DownloadToS3Block(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
# get workflow run context
@@ -1830,7 +1822,6 @@ class UploadToS3Block(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
# get workflow run context
@@ -2008,7 +1999,6 @@ class FileUploadBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
# get workflow run context
@@ -2426,7 +2416,6 @@ class SendEmailBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
workflow_run_context = self.get_workflow_run_context(workflow_run_id)
@@ -2662,7 +2651,6 @@ class FileParserBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
workflow_run_context = self.get_workflow_run_context(workflow_run_id)
@@ -2802,7 +2790,6 @@ class PDFParserBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
workflow_run_context = self.get_workflow_run_context(workflow_run_id)
@@ -2903,7 +2890,6 @@ class WaitBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
# TODO: we need to support to interrupt the sleep when the workflow run failed/cancelled/terminated
@@ -2948,7 +2934,6 @@ class ValidationBlock(BaseTaskBlock):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
task_order, _ = await self.get_task_order(workflow_run_id, 0)
@@ -3051,7 +3036,6 @@ class TaskV2Block(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunStatus # noqa: PLC0415
@@ -3123,7 +3107,6 @@ class TaskV2Block(Block):
request_id=None,
max_steps_override=self.max_steps,
browser_session_id=browser_session_id,
code_gen=code_gen,
)
finally:
context: skyvern_context.SkyvernContext | None = skyvern_context.current()
@@ -3241,7 +3224,6 @@ class HttpRequestBlock(Block):
workflow_run_block_id: str,
organization_id: str | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
**kwargs: dict,
) -> BlockResult:
"""Execute the HTTP request and return the response"""

View File

@@ -145,6 +145,7 @@ class WorkflowRun(BaseModel):
job_id: str | None = None
sequential_key: str | None = None
ai_fallback: bool | None = None
code_gen: bool | None = None
queued_at: datetime | None = None
started_at: datetime | None = None

View File

@@ -197,6 +197,7 @@ class WorkflowService:
max_steps_override: int | None = None,
parent_workflow_run_id: str | None = None,
debug_session_id: str | None = None,
code_gen: bool | None = None,
) -> WorkflowRun:
"""
Create a workflow run and its parameters. Validate the workflow and the organization. If there are missing
@@ -231,6 +232,7 @@ class WorkflowService:
parent_workflow_run_id=parent_workflow_run_id,
sequential_key=workflow.sequential_key,
debug_session_id=debug_session_id,
code_gen=code_gen,
)
LOG.info(
f"Created workflow run {workflow_run.workflow_run_id} for workflow {workflow.workflow_id}",
@@ -243,6 +245,7 @@ class WorkflowService:
max_screenshot_scrolling_times=workflow_request.max_screenshot_scrolls,
ai_fallback=workflow_request.ai_fallback,
run_with=workflow_request.run_with,
code_gen=code_gen,
)
context: skyvern_context.SkyvernContext | None = skyvern_context.current()
current_run_id = context.run_id if context and context.run_id else workflow_run.workflow_run_id
@@ -322,7 +325,6 @@ class WorkflowService:
block_labels: list[str] | None = None,
block_outputs: dict[str, Any] | None = None,
browser_session_id: str | None = None,
code_gen: bool | None = None,
) -> WorkflowRun:
"""Execute a workflow."""
organization_id = organization.organization_id
@@ -407,8 +409,11 @@ class WorkflowService:
# Check if there's a related workflow script that should be used instead
workflow_script, _ = await workflow_script_service.get_workflow_script(workflow, workflow_run, block_labels)
current_context = skyvern_context.current()
if workflow_script and current_context:
current_context.published_workflow_script_id = workflow_script.script_id
if current_context:
if workflow_script:
current_context.generate_script = False
if workflow_run.code_gen:
current_context.generate_script = True
is_script_run = self.should_run_script(workflow, workflow_run)
if workflow_script and is_script_run:
LOG.info(
@@ -452,7 +457,6 @@ class WorkflowService:
workflow=workflow,
workflow_run=workflow_run,
block_labels=block_labels,
code_gen=code_gen,
)
else:
LOG.info(
@@ -1053,6 +1057,7 @@ class WorkflowService:
parent_workflow_run_id: str | None = None,
sequential_key: str | None = None,
debug_session_id: str | None = None,
code_gen: bool | None = None,
) -> WorkflowRun:
# validate the browser session id
if workflow_request.browser_session_id:
@@ -1080,6 +1085,7 @@ class WorkflowService:
run_with=workflow_request.run_with,
debug_session_id=debug_session_id,
ai_fallback=workflow_request.ai_fallback,
code_gen=code_gen,
)
async def _update_workflow_run_status(
@@ -2684,8 +2690,8 @@ class WorkflowService:
workflow: Workflow,
workflow_run: WorkflowRun,
block_labels: list[str] | None = None,
code_gen: bool | None = None,
) -> None:
code_gen = workflow_run.code_gen
LOG.info(
"Generate script?",
block_labels=block_labels,