diff --git a/alembic/versions/2024_12_31_0936-521241e64aed_re_introduce_indexes_for_artifacts.py b/alembic/versions/2024_12_31_0936-521241e64aed_re_introduce_indexes_for_artifacts.py new file mode 100644 index 00000000..ea0d26ca --- /dev/null +++ b/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 ### diff --git a/skyvern/forge/sdk/api/llm/models.py b/skyvern/forge/sdk/api/llm/models.py index c0a920c9..c913cc0c 100644 --- a/skyvern/forge/sdk/api/llm/models.py +++ b/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]]: ... diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 7c1b0871..04bf968e 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -161,20 +161,14 @@ class OrganizationAuthTokenModel(Base): class ArtifactModel(Base): __tablename__ = "artifacts" - __table_args__ = ( - 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"), - ) + __table_args__ = (Index("org_task_step_index", "organization_id", "task_id", "step_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_block_id = Column(String) - observer_cruise_id = Column(String) - observer_thought_id = Column(String) + workflow_run_id = Column(String, index=True) + workflow_run_block_id = Column(String, index=True) + observer_cruise_id = Column(String, index=True) + 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)