Skip to content

Commit

Permalink
Make text prompt executions a step in workflow runs (#1590)
Browse files Browse the repository at this point in the history
  • Loading branch information
wintonzheng authored Jan 17, 2025
1 parent bb834fa commit e4cb098
Show file tree
Hide file tree
Showing 4 changed files with 25 additions and 1 deletion.
7 changes: 7 additions & 0 deletions skyvern/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -311,6 +311,13 @@ def __init__(self, step_id: str, reason: str) -> None:
super().__init__(f"Step {step_id} cannot be executed and task is failed. Reason: {reason}")


class BlockTerminationError(SkyvernException):
def __init__(self, workflow_run_block_id: str, workflow_run_id: str, reason: str) -> None:
super().__init__(
f"Block {workflow_run_block_id} cannot be executed and workflow run {workflow_run_id} is failed. Reason: {reason}"
)


class StepUnableToExecuteError(SkyvernException):
def __init__(self, step_id: str, reason: str) -> None:
super().__init__(f"Step {step_id} cannot be executed and task execution is stopped. Reason: {reason}")
Expand Down
6 changes: 6 additions & 0 deletions skyvern/forge/agent_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from skyvern.forge.sdk.models import Step, StepStatus
from skyvern.forge.sdk.schemas.organizations import Organization
from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus
from skyvern.forge.sdk.workflow.models.block import BlockTypeVar
from skyvern.webeye.browser_factory import BrowserState
from skyvern.webeye.scraper.scraper import ELEMENT_NODE_ATTRIBUTES, CleanupElementTreeFunc, json_to_html

Expand Down Expand Up @@ -379,6 +380,11 @@ async def validate_step_execution(
if not can_execute:
raise StepUnableToExecuteError(step_id=step.step_id, reason=f"Cannot execute step. Reasons: {reasons}")

async def validate_block_execution(
self, block: BlockTypeVar, workflow_run_id: str, workflow_run_block_id: str, organization_id: str | None
) -> None:
return

async def prepare_step_execution(
self,
organization: Organization | None,
Expand Down
7 changes: 7 additions & 0 deletions skyvern/forge/sdk/workflow/models/block.py
Original file line number Diff line number Diff line change
Expand Up @@ -1181,6 +1181,13 @@ async def execute(
browser_session_id: str | None = None,
**kwargs: dict,
) -> BlockResult:
# Validate block execution
await app.AGENT_FUNCTION.validate_block_execution(
block=self,
workflow_run_block_id=workflow_run_block_id,
workflow_run_id=workflow_run_id,
organization_id=organization_id,
)
# get workflow run context
workflow_run_context = self.get_workflow_run_context(workflow_run_id)
await app.DATABASE.update_workflow_run_block(
Expand Down
6 changes: 5 additions & 1 deletion skyvern/forge/sdk/workflow/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -860,11 +860,15 @@ async def build_workflow_run_status_response(
workflow_run_steps = await app.DATABASE.get_steps_by_task_ids(
task_ids=[task.task_id for task in workflow_run_tasks], organization_id=organization_id
)
workflow_run_blocks = await app.DATABASE.get_workflow_run_blocks(
workflow_run_id=workflow_run_id, organization_id=organization_id
)
text_prompt_blocks = [block for block in workflow_run_blocks if block.block_type == BlockType.TEXT_PROMPT]
total_steps = len(workflow_run_steps)
# TODO: This is a temporary cost calculation. We need to implement a more accurate cost calculation.
# successful steps are the ones that have a status of completed and the total count of unique step.order
successful_steps = [step for step in workflow_run_steps if step.status == StepStatus.completed]
total_cost = 0.1 * len(successful_steps)
total_cost = 0.1 * (len(successful_steps) + len(text_prompt_blocks))
return WorkflowRunStatusResponse(
workflow_id=workflow.workflow_permanent_id,
workflow_run_id=workflow_run_id,
Expand Down

0 comments on commit e4cb098

Please sign in to comment.