diff --git a/alembic/versions/2025_08_28_2112-d3ec63728c2a_add_ai_fallback_field_to_workflows_table.py b/alembic/versions/2025_08_28_2112-d3ec63728c2a_add_ai_fallback_field_to_workflows_table.py new file mode 100644 index 00000000..df5c8b8d --- /dev/null +++ b/alembic/versions/2025_08_28_2112-d3ec63728c2a_add_ai_fallback_field_to_workflows_table.py @@ -0,0 +1,31 @@ +"""add ai_fallback field to workflows table + +Revision ID: d3ec63728c2a +Revises: 1bba8a38ddc7 +Create Date: 2025-08-28 21:12:54.750395+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "d3ec63728c2a" +down_revision: Union[str, None] = "1bba8a38ddc7" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("workflows", sa.Column("ai_fallback", sa.Boolean(), nullable=False, server_default=sa.false())) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("workflows", "ai_fallback") + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 425ebc7b..2e845e3f 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/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: diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index e2885217..1fa1c1aa 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/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) diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index f507c7fc..7bbe4501 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/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, ) diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py index 8cb2470b..05d9d881 100644 --- a/skyvern/forge/sdk/workflow/models/workflow.py +++ b/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 diff --git a/skyvern/schemas/workflows.py b/skyvern/schemas/workflows.py index b63f4ce1..8721218b 100644 --- a/skyvern/schemas/workflows.py +++ b/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 diff --git a/skyvern/services/script_service.py b/skyvern/services/script_service.py index bb2335b3..2af0e648 100644 --- a/skyvern/services/script_service.py +++ b/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,