recreate index for artifacts table (#1455)

alembic/versions/2024_12_31_0936-521241e64aed_re_introduce_indexes_for_artifacts.py

@@ -0,0 +1,43 @@
+"""re-introduce indexes for artifacts
+
+Revision ID: 521241e64aed
+Revises: 172cdfb3e2ee
+Create Date: 2024-12-31 09:36:18.974818+00:00
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "521241e64aed"
+down_revision: Union[str, None] = "172cdfb3e2ee"
+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.drop_index("org_observer_cruise_index", table_name="artifacts")
+    op.drop_index("org_observer_thought_index", table_name="artifacts")
+    op.drop_index("org_wfrb_index", table_name="artifacts")
+    op.drop_index("org_workflow_run_index", table_name="artifacts")
+    op.create_index(op.f("ix_artifacts_observer_cruise_id"), "artifacts", ["observer_cruise_id"], unique=False)
+    op.create_index(op.f("ix_artifacts_observer_thought_id"), "artifacts", ["observer_thought_id"], unique=False)
+    op.create_index(op.f("ix_artifacts_workflow_run_block_id"), "artifacts", ["workflow_run_block_id"], unique=False)
+    op.create_index(op.f("ix_artifacts_workflow_run_id"), "artifacts", ["workflow_run_id"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_artifacts_workflow_run_id"), table_name="artifacts")
+    op.drop_index(op.f("ix_artifacts_workflow_run_block_id"), table_name="artifacts")
+    op.drop_index(op.f("ix_artifacts_observer_thought_id"), table_name="artifacts")
+    op.drop_index(op.f("ix_artifacts_observer_cruise_id"), table_name="artifacts")
+    op.create_index("org_workflow_run_index", "artifacts", ["organization_id", "workflow_run_id"], unique=False)
+    op.create_index("org_wfrb_index", "artifacts", ["organization_id", "workflow_run_block_id"], unique=False)
+    op.create_index("org_observer_thought_index", "artifacts", ["organization_id", "observer_thought_id"], unique=False)
+    op.create_index("org_observer_cruise_index", "artifacts", ["organization_id", "observer_cruise_id"], unique=False)
+    # ### end Alembic commands ###

skyvern/forge/sdk/api/llm/models.py

@@ -4,7 +4,7 @@ from typing import Any, Awaitable, Literal, Optional, Protocol, TypedDict
 from litellm import AllowedFailsPolicy
 
 from skyvern.forge.sdk.models import Step
-from skyvern.forge.sdk.schemas.observers import ObserverCruise
+from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
 from skyvern.forge.sdk.settings_manager import SettingsManager
 
 
@@ -80,7 +80,7 @@ class LLMAPIHandler(Protocol):
         prompt: str,
         step: Step | None = None,
         observer_cruise: ObserverCruise | None = None,
-        observer_thought: ObserverCruise | None = None,
+        observer_thought: ObserverThought | None = None,
         screenshots: list[bytes] | None = None,
         parameters: dict[str, Any] | None = None,
     ) -> Awaitable[dict[str, Any]]: ...

skyvern/forge/sdk/db/models.py

@@ -161,20 +161,14 @@ class OrganizationAuthTokenModel(Base):
 
 class ArtifactModel(Base):
     __tablename__ = "artifacts"
-    __table_args__ = (
+    __table_args__ = (Index("org_task_step_index", "organization_id", "task_id", "step_id"),)
-        Index("org_task_step_index", "organization_id", "task_id", "step_id"),
-        Index("org_workflow_run_index", "organization_id", "workflow_run_id"),
-        Index("org_observer_cruise_index", "organization_id", "observer_cruise_id"),
-        Index("org_observer_thought_index", "organization_id", "observer_thought_id"),
-        Index("org_wfrb_index", "organization_id", "workflow_run_block_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)
+    workflow_run_id = Column(String, index=True)
-    workflow_run_block_id = Column(String)
+    workflow_run_block_id = Column(String, index=True)
-    observer_cruise_id = Column(String)
+    observer_cruise_id = Column(String, index=True)
-    observer_thought_id = Column(String)
+    observer_thought_id = Column(String, index=True)
     task_id = Column(String, ForeignKey("tasks.task_id"))
     step_id = Column(String, ForeignKey("steps.step_id"), index=True)
     artifact_type = Column(String)