From 944c95f4561ac5fb13258e7bd2df3f109fec82b7 Mon Sep 17 00:00:00 2001 From: Marc Kelechava Date: Thu, 4 Dec 2025 13:47:46 -0800 Subject: [PATCH] add cached steps flag to db (is_script_cached in Step) (#4202) --- ...354699b93_db_migration_created_by_steps.py | 31 +++++++++++++++++++ skyvern/forge/sdk/db/client.py | 5 +++ skyvern/forge/sdk/db/models.py | 1 + skyvern/forge/sdk/db/utils.py | 1 + skyvern/forge/sdk/models.py | 1 + skyvern/services/script_service.py | 7 +++++ 6 files changed, 46 insertions(+) create mode 100644 alembic/versions/2025_12_04_2120-152354699b93_db_migration_created_by_steps.py diff --git a/alembic/versions/2025_12_04_2120-152354699b93_db_migration_created_by_steps.py b/alembic/versions/2025_12_04_2120-152354699b93_db_migration_created_by_steps.py new file mode 100644 index 00000000..a1b44147 --- /dev/null +++ b/alembic/versions/2025_12_04_2120-152354699b93_db_migration_created_by_steps.py @@ -0,0 +1,31 @@ +"""db migration created_by steps + +Revision ID: 152354699b93 +Revises: 44e95cb21d98 +Create Date: 2025-12-04 21:20:20.166647+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "152354699b93" +down_revision: Union[str, None] = "44e95cb21d98" +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("steps", sa.Column("created_by", sa.String(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("steps", "created_by") + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 59218f1b..6fead0d7 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -257,6 +257,7 @@ class AgentDB: retry_index: int, organization_id: str | None = None, status: StepStatus = StepStatus.created, + created_by: str | None = None, ) -> Step: try: async with self.Session() as session: @@ -266,6 +267,7 @@ class AgentDB: retry_index=retry_index, status=status, organization_id=organization_id, + created_by=created_by, ) session.add(new_step) await session.commit() @@ -595,6 +597,7 @@ class AgentDB: incremental_output_tokens: int | None = None, incremental_reasoning_tokens: int | None = None, incremental_cached_tokens: int | None = None, + created_by: str | None = None, ) -> Step: try: async with self.Session() as session: @@ -627,6 +630,8 @@ class AgentDB: step.reasoning_token_count = incremental_reasoning_tokens + (step.reasoning_token_count or 0) if incremental_cached_tokens is not None: step.cached_token_count = incremental_cached_tokens + (step.cached_token_count or 0) + if created_by is not None: + step.created_by = created_by await session.commit() updated_step = await self.get_step(step_id, organization_id) diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 9aa253c4..1fd5070c 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -142,6 +142,7 @@ class StepModel(Base): cached_token_count = Column(Integer, default=0) step_cost = Column(Numeric, default=0) finished_at = Column(DateTime, nullable=True) + created_by = Column(String, nullable=True) class OrganizationModel(Base): diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index baf9de2b..f37d877c 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -234,6 +234,7 @@ def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step: reasoning_token_count=step_model.reasoning_token_count, cached_token_count=step_model.cached_token_count, step_cost=step_model.step_cost, + created_by=step_model.created_by, ) diff --git a/skyvern/forge/sdk/models.py b/skyvern/forge/sdk/models.py index de0730ed..86ef0457 100644 --- a/skyvern/forge/sdk/models.py +++ b/skyvern/forge/sdk/models.py @@ -71,6 +71,7 @@ class Step(BaseModel): reasoning_token_count: int | None = None cached_token_count: int | None = None step_cost: float = 0 + created_by: str | None = None is_speculative: bool = False speculative_original_status: StepStatus | None = None speculative_llm_metadata: SpeculativeLLMMetadata | None = None diff --git a/skyvern/services/script_service.py b/skyvern/services/script_service.py index 4332344d..8669331c 100644 --- a/skyvern/services/script_service.py +++ b/skyvern/services/script_service.py @@ -398,6 +398,7 @@ async def _create_workflow_block_run_and_task( url: str | None = None, label: str | None = None, model: dict[str, Any] | None = None, + created_by: str | None = None, ) -> tuple[str | None, str | None, str | None]: """ Create a workflow block run and optionally a task if workflow_run_id is available in context. @@ -460,6 +461,7 @@ async def _create_workflow_block_run_and_task( retry_index=0, organization_id=organization_id, status=StepStatus.running, + created_by=created_by, ) step_id = step.step_id # reset the action order to 0 @@ -1334,6 +1336,7 @@ async def run_task( url=url, label=cache_key, model=model, + created_by="script", ) prompt = _render_template_with_label(prompt, cache_key) # set the prompt in the RunContext @@ -1419,6 +1422,7 @@ async def download( url=url, label=cache_key, model=model, + created_by="script", ) prompt = _render_template_with_label(prompt, cache_key) # set the prompt in the RunContext @@ -1499,6 +1503,7 @@ async def action( url=url, label=cache_key, model=model, + created_by="script", ) prompt = _render_template_with_label(prompt, cache_key) # set the prompt in the RunContext @@ -1578,6 +1583,7 @@ async def login( url=url, label=cache_key, model=model, + created_by="script", ) prompt = _render_template_with_label(prompt, cache_key) if totp_url: @@ -1660,6 +1666,7 @@ async def extract( url=url, label=cache_key, model=model, + created_by="script", ) prompt = _render_template_with_label(prompt, cache_key) # set the prompt in the RunContext