From e4cb0987cb0dfe1097e94cc74a6f3034e8747585 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Fri, 17 Jan 2025 09:17:31 -0800
Subject: [PATCH] Make text prompt executions a step in workflow runs (#1590)

---
 skyvern/exceptions.py                      | 7 +++++++
 skyvern/forge/agent_functions.py           | 6 ++++++
 skyvern/forge/sdk/workflow/models/block.py | 7 +++++++
 skyvern/forge/sdk/workflow/service.py      | 6 +++++-
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py
index 16be5eba..24da926d 100644
--- a/skyvern/exceptions.py
+++ b/skyvern/exceptions.py
@@ -311,6 +311,13 @@ class StepTerminationError(SkyvernException):
         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}")
diff --git a/skyvern/forge/agent_functions.py b/skyvern/forge/agent_functions.py
index a0d8740e..e02cfec7 100644
--- a/skyvern/forge/agent_functions.py
+++ b/skyvern/forge/agent_functions.py
@@ -17,6 +17,7 @@ from skyvern.forge.sdk.api.llm.exceptions import LLMProviderError
 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
 
@@ -379,6 +380,11 @@ class AgentFunction:
         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,
diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py
index a6bc23b3..b237ca1d 100644
--- a/skyvern/forge/sdk/workflow/models/block.py
+++ b/skyvern/forge/sdk/workflow/models/block.py
@@ -1181,6 +1181,13 @@ class TextPromptBlock(Block):
         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(
diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index 00a763ad..96f9c1be 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -860,11 +860,15 @@ class WorkflowService:
             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,