From f1733a50546b40c2303f31e7e2ed7f720bd2ec67 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Thu, 5 Dec 2024 23:16:41 -0800 Subject: [PATCH] add workflow_run_id column to artifacts + ObserverCruise and ObserverThought (#1298) --- ...96d9557da_introduce_observercruise_and_.py | 91 +++++++++++++++++++ skyvern/forge/sdk/db/id.py | 12 +++ skyvern/forge/sdk/db/models.py | 31 ++++++- 3 files changed, 133 insertions(+), 1 deletion(-) create mode 100644 alembic/versions/2024_12_06_0540-3c196d9557da_introduce_observercruise_and_.py diff --git a/alembic/versions/2024_12_06_0540-3c196d9557da_introduce_observercruise_and_.py b/alembic/versions/2024_12_06_0540-3c196d9557da_introduce_observercruise_and_.py new file mode 100644 index 00000000..53ea9a03 --- /dev/null +++ b/alembic/versions/2024_12_06_0540-3c196d9557da_introduce_observercruise_and_.py @@ -0,0 +1,91 @@ +"""Introduce ObserverCruise and ObserverThought. Add workflow_run_block_id to artifacts + +Revision ID: 3c196d9557da +Revises: de0254717601 +Create Date: 2024-12-06 05:40:10.408358+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "3c196d9557da" +down_revision: Union[str, None] = "de0254717601" +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.create_table( + "observer_cruises", + sa.Column("observer_cruise_id", sa.String(), nullable=False), + sa.Column("status", sa.String(), nullable=False), + sa.Column("organization_id", sa.String(), nullable=True), + sa.Column("workflow_run_id", sa.String(), nullable=True), + sa.Column("workflow_id", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["organization_id"], + ["organizations.organization_id"], + ), + sa.ForeignKeyConstraint( + ["workflow_id"], + ["workflows.workflow_id"], + ), + sa.ForeignKeyConstraint( + ["workflow_run_id"], + ["workflow_runs.workflow_run_id"], + ), + sa.PrimaryKeyConstraint("observer_cruise_id"), + ) + op.create_table( + "observer_thoughts", + sa.Column("observer_thought_id", sa.String(), nullable=False), + sa.Column("organization_id", sa.String(), nullable=True), + sa.Column("observer_cruise_id", sa.String(), nullable=False), + sa.Column("workflow_run_id", sa.String(), nullable=True), + sa.Column("workflow_id", sa.String(), nullable=True), + sa.Column("thought", sa.String(), nullable=True), + sa.Column("answer", sa.String(), nullable=True), + sa.ForeignKeyConstraint( + ["observer_cruise_id"], + ["observer_cruises.observer_cruise_id"], + ), + sa.ForeignKeyConstraint( + ["organization_id"], + ["organizations.organization_id"], + ), + sa.ForeignKeyConstraint( + ["workflow_id"], + ["workflows.workflow_id"], + ), + sa.ForeignKeyConstraint( + ["workflow_run_id"], + ["workflow_runs.workflow_run_id"], + ), + sa.PrimaryKeyConstraint("observer_thought_id"), + ) + op.add_column("artifacts", sa.Column("workflow_run_id", sa.String(), nullable=True)) + op.add_column("artifacts", sa.Column("workflow_run_block_id", sa.String(), nullable=True)) + op.create_index("org_workflow_run_index", "artifacts", ["organization_id", "workflow_run_id"], unique=False) + op.create_foreign_key(None, "artifacts", "workflow_runs", ["workflow_run_id"], ["workflow_run_id"]) + op.create_foreign_key( + None, "artifacts", "workflow_run_blocks", ["workflow_run_block_id"], ["workflow_run_block_id"] + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_constraint(None, "artifacts", type_="foreignkey") + op.drop_constraint(None, "artifacts", type_="foreignkey") + op.drop_index("org_workflow_run_index", table_name="artifacts") + op.drop_column("artifacts", "workflow_run_block_id") + op.drop_column("artifacts", "workflow_run_id") + op.drop_table("observer_thoughts") + op.drop_table("observer_cruises") + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/id.py b/skyvern/forge/sdk/db/id.py index da992cf5..a27242ab 100644 --- a/skyvern/forge/sdk/db/id.py +++ b/skyvern/forge/sdk/db/id.py @@ -44,6 +44,8 @@ BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc" BITWARDEN_SENSITIVE_INFORMATION_PARAMETER_PREFIX = "bsi" BITWARDEN_CREDIT_CARD_DATA_PARAMETER_PREFIX = "bccd" TASK_GENERATION_PREFIX = "tg" +OBSERVER_CRUISE_ID = "oc" +OBSERVER_THOUGHT_ID = "ot" def generate_workflow_id() -> str: @@ -141,6 +143,16 @@ def generate_action_id() -> str: return f"a_{int_id}" +def generate_observer_cruise_id() -> str: + int_id = generate_id() + return f"{OBSERVER_CRUISE_ID}_{int_id}" + + +def generate_observer_thought_id() -> str: + int_id = generate_id() + return f"{OBSERVER_THOUGHT_ID}_{int_id}" + + def generate_id() -> int: """ generate a 64-bit int ID diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 6c61db95..4dd2b2ef 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -25,6 +25,8 @@ from skyvern.forge.sdk.db.id import ( generate_bitwarden_credit_card_data_parameter_id, generate_bitwarden_login_credential_parameter_id, generate_bitwarden_sensitive_information_parameter_id, + generate_observer_cruise_id, + generate_observer_thought_id, generate_org_id, generate_organization_auth_token_id, generate_output_parameter_id, @@ -157,10 +159,15 @@ class OrganizationAuthTokenModel(Base): class ArtifactModel(Base): __tablename__ = "artifacts" - __table_args__ = (Index("org_task_step_index", "organization_id", "task_id", "step_id"),) + __table_args__ = ( + Index("org_task_step_index", "organization_id", "task_id", "step_id"), + Index("org_workflow_run_index", "organization_id", "workflow_run_id"), + ) artifact_id = Column(String, primary_key=True, index=True, default=generate_artifact_id) organization_id = Column(String, ForeignKey("organizations.organization_id")) + workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id")) + workflow_run_block_id = Column(String, ForeignKey("workflow_run_blocks.workflow_run_block_id")) task_id = Column(String, ForeignKey("tasks.task_id")) step_id = Column(String, ForeignKey("steps.step_id"), index=True) artifact_type = Column(String) @@ -495,3 +502,25 @@ class WorkflowRunBlockModel(Base): created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False) + + +class ObserverCruise(Base): + __tablename__ = "observer_cruises" + + observer_cruise_id = Column(String, primary_key=True, default=generate_observer_cruise_id) + status = Column(String, nullable=False, default="created") + organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True) + workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True) + workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=True) + + +class ObserverThought(Base): + __tablename__ = "observer_thoughts" + + observer_thought_id = Column(String, primary_key=True, default=generate_observer_thought_id) + organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True) + observer_cruise_id = Column(String, ForeignKey("observer_cruises.observer_cruise_id"), nullable=False) + workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True) + workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=True) + thought = Column(String, nullable=True) + answer = Column(String, nullable=True)