flag to control ai fallback (#3313)

skyvern/forge/sdk/db/client.py

@@ -1366,6 +1366,7 @@ class AgentDB:
         is_saved_task: bool = False,
         status: WorkflowStatus = WorkflowStatus.published,
         generate_script: bool = False,
+        ai_fallback: bool = False,
         cache_key: str | None = None,
     ) -> Workflow:
         async with self.Session() as session:
@@ -1385,6 +1386,7 @@ class AgentDB:
                 is_saved_task=is_saved_task,
                 status=status,
                 generate_script=generate_script,
+                ai_fallback=ai_fallback,
                 cache_key=cache_key,
             )
             if workflow_permanent_id:

skyvern/forge/sdk/db/models.py

@@ -243,6 +243,7 @@ class WorkflowModel(Base):
     model = Column(JSON, nullable=True)
     status = Column(String, nullable=False, default="published")
     generate_script = Column(Boolean, default=False, nullable=False)
+    ai_fallback = Column(Boolean, default=False, nullable=False)
     cache_key = Column(String, nullable=True)
 
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)

skyvern/forge/sdk/db/utils.py

@@ -264,6 +264,7 @@ def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = Fal
         status=WorkflowStatus(workflow_model.status),
         extra_http_headers=workflow_model.extra_http_headers,
         generate_script=workflow_model.generate_script,
+        ai_fallback=workflow_model.ai_fallback,
         cache_key=workflow_model.cache_key,
     )
 

skyvern/forge/sdk/workflow/models/workflow.py

@@ -77,6 +77,7 @@ class Workflow(BaseModel):
     max_screenshot_scrolls: int | None = None
     extra_http_headers: dict[str, str] | None = None
     generate_script: bool = False
+    ai_fallback: bool = False
     cache_key: str | None = None
 
     created_at: datetime

skyvern/schemas/workflows.py

@@ -506,6 +506,7 @@ class WorkflowCreateYAMLRequest(BaseModel):
     extra_http_headers: dict[str, str] | None = None
     status: WorkflowStatus = WorkflowStatus.published
     generate_script: bool = False
+    ai_fallback: bool = False
     cache_key: str | None = None
 
 

skyvern/services/script_service.py

@@ -478,7 +478,7 @@ async def _fallback_to_ai_run(
     try:
         organization_id = context.organization_id
         LOG.info(
-            "Script fallback to AI run",
+            "Script trying to fallback to AI run",
             cache_key=cache_key,
             organization_id=organization_id,
             workflow_id=context.workflow_id,
@@ -510,13 +510,22 @@ async def _fallback_to_ai_run(
         if not task:
             raise Exception(f"Task is missing task_id={context.task_id}")
         workflow = await app.DATABASE.get_workflow(workflow_id=context.workflow_id, organization_id=organization_id)
-        if not workflow:
+        if not workflow or not workflow.ai_fallback:
             return
 
         # get the output_paramter
         output_parameter = workflow.get_output_parameter(cache_key)
         if not output_parameter:
             return
+        LOG.info(
+            "Script starting to fallback to AI run",
+            cache_key=cache_key,
+            organization_id=organization_id,
+            workflow_id=context.workflow_id,
+            workflow_run_id=context.workflow_run_id,
+            task_id=context.task_id,
+            step_id=context.step_id,
+        )
 
         task_block = TaskBlock(
             label=cache_key,