add workflow_run_id column to artifacts + ObserverCruise and ObserverThought (#1298)

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 ###

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

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)