adopt ruff as the replacement for python black (#332)

2024-05-16 18:20:11 -07:00
parent 7a2be7e355
commit 2466897158
44 changed files with 1081 additions and 321 deletions
--- a/skyvern/forge/sdk/db/client.py
+++ b/skyvern/forge/sdk/db/client.py
@@ -184,7 +184,11 @@ class AgentDB:
                ).first():
                    return convert_to_task(task_obj, self.debug_enabled)
                else:
-                    LOG.info("Task not found", task_id=task_id, organization_id=organization_id)
+                    LOG.info(
+                        "Task not found",
+                        task_id=task_id,
+                        organization_id=organization_id,
+                    )
                    return None
        except SQLAlchemyError:
            LOG.error("SQLAlchemyError", exc_info=True)
@@ -266,7 +270,11 @@ class AgentDB:
                ).first():
                    return convert_to_step(step, debug_enabled=self.debug_enabled)
                else:
-                    LOG.info("Latest step not found", task_id=task_id, organization_id=organization_id)
+                    LOG.info(
+                        "Latest step not found",
+                        task_id=task_id,
+                        organization_id=organization_id,
+                    )
                    return None
        except SQLAlchemyError:
            LOG.error("SQLAlchemyError", exc_info=True)
@@ -812,7 +820,10 @@ class AgentDB:
                    )
                    .where(WorkflowModel.organization_id == organization_id)
                    .where(WorkflowModel.deleted_at.is_(None))
-                    .group_by(WorkflowModel.organization_id, WorkflowModel.workflow_permanent_id)
+                    .group_by(
+                        WorkflowModel.organization_id,
+                        WorkflowModel.workflow_permanent_id,
+                    )
                    .subquery()
                )
                main_query = (
@@ -924,7 +935,10 @@ class AgentDB:
                await session.commit()
                await session.refresh(workflow_run)
                return convert_to_workflow_run(workflow_run)
-            LOG.error("WorkflowRun not found, nothing to update", workflow_run_id=workflow_run_id)
+            LOG.error(
+                "WorkflowRun not found, nothing to update",
+                workflow_run_id=workflow_run_id,
+            )
            return None

    async def get_workflow_run(self, workflow_run_id: str) -> WorkflowRun | None:
@@ -1066,7 +1080,10 @@ class AgentDB:
            raise

    async def create_workflow_run_output_parameter(
-        self, workflow_run_id: str, output_parameter_id: str, value: dict[str, Any] | list | str | None
+        self,
+        workflow_run_id: str,
+        output_parameter_id: str,
+        value: dict[str, Any] | list | str | None,
    ) -> WorkflowRunOutputParameter:
        try:
            async with self.Session() as session:
@@ -1149,7 +1166,9 @@ class AgentDB:
                        (
                            workflow_parameter,
                            convert_to_workflow_run_parameter(
-                                workflow_run_parameter, workflow_parameter, self.debug_enabled
+                                workflow_run_parameter,
+                                workflow_parameter,
+                                self.debug_enabled,
                            ),
                        )
                    )
--- a/skyvern/forge/sdk/db/models.py
+++ b/skyvern/forge/sdk/db/models.py
@@ -63,7 +63,11 @@ class TaskModel(Base):
    max_steps_per_run = Column(Integer, nullable=True)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False, index=True)
    modified_at = Column(
-        DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False, index=True
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+        index=True,
    )


@@ -80,7 +84,12 @@ class StepModel(Base):
    is_last = Column(Boolean, default=False)
    retry_index = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )
    input_token_count = Column(Integer, default=0)
    output_token_count = Column(Integer, default=0)
    step_cost = Column(Numeric, default=0)
@@ -96,7 +105,12 @@ class OrganizationModel(Base):
    max_retries_per_step = Column(Integer, nullable=True)
    domain = Column(String, nullable=True, index=True)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime,
+        nullable=False,
+    )


 class OrganizationAuthTokenModel(Base):
@@ -115,7 +129,12 @@ class OrganizationAuthTokenModel(Base):
    valid = Column(Boolean, nullable=False, default=True)

    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime,
+        nullable=False,
+    )
    deleted_at = Column(DateTime, nullable=True)


@@ -130,13 +149,23 @@ class ArtifactModel(Base):
    artifact_type = Column(String)
    uri = Column(String)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )


 class WorkflowModel(Base):
    __tablename__ = "workflows"
    __table_args__ = (
-        UniqueConstraint("organization_id", "workflow_permanent_id", "version", name="uc_org_permanent_id_version"),
+        UniqueConstraint(
+            "organization_id",
+            "workflow_permanent_id",
+            "version",
+            name="uc_org_permanent_id_version",
+        ),
        Index("permanent_id_version_idx", "workflow_permanent_id", "version"),
    )

@@ -149,7 +178,12 @@ class WorkflowModel(Base):
    webhook_callback_url = Column(String)

    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )
    deleted_at = Column(DateTime, nullable=True)

    workflow_permanent_id = Column(String, nullable=False, default=generate_workflow_permanent_id, index=True)
@@ -166,7 +200,12 @@ class WorkflowRunModel(Base):
    webhook_callback_url = Column(String)

    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )


 class WorkflowParameterModel(Base):
@@ -179,7 +218,12 @@ class WorkflowParameterModel(Base):
    workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False)
    default_value = Column(String, nullable=True)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )
    deleted_at = Column(DateTime, nullable=True)


@@ -191,7 +235,12 @@ class OutputParameterModel(Base):
    description = Column(String, nullable=True)
    workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )
    deleted_at = Column(DateTime, nullable=True)


@@ -204,7 +253,12 @@ class AWSSecretParameterModel(Base):
    description = Column(String, nullable=True)
    aws_key = Column(String, nullable=False)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )
    deleted_at = Column(DateTime, nullable=True)


@@ -212,7 +266,10 @@ class BitwardenLoginCredentialParameterModel(Base):
    __tablename__ = "bitwarden_login_credential_parameters"

    bitwarden_login_credential_parameter_id = Column(
-        String, primary_key=True, index=True, default=generate_bitwarden_login_credential_parameter_id
+        String,
+        primary_key=True,
+        index=True,
+        default=generate_bitwarden_login_credential_parameter_id,
    )
    workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False)
    key = Column(String, nullable=False)
@@ -222,16 +279,29 @@ class BitwardenLoginCredentialParameterModel(Base):
    bitwarden_master_password_aws_secret_key = Column(String, nullable=False)
    url_parameter_key = Column(String, nullable=False)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime,
+        default=datetime.datetime.utcnow,
+        onupdate=datetime.datetime.utcnow,
+        nullable=False,
+    )
    deleted_at = Column(DateTime, nullable=True)


 class WorkflowRunParameterModel(Base):
    __tablename__ = "workflow_run_parameters"

-    workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), primary_key=True, index=True)
+    workflow_run_id = Column(
+        String,
+        ForeignKey("workflow_runs.workflow_run_id"),
+        primary_key=True,
+        index=True,
+    )
    workflow_parameter_id = Column(
-        String, ForeignKey("workflow_parameters.workflow_parameter_id"), primary_key=True, index=True
+        String,
+        ForeignKey("workflow_parameters.workflow_parameter_id"),
+        primary_key=True,
+        index=True,
    )
    # Can be bool | int | float | str | dict | list depending on the workflow parameter type
    value = Column(String, nullable=False)
@@ -241,9 +311,17 @@ class WorkflowRunParameterModel(Base):
 class WorkflowRunOutputParameterModel(Base):
    __tablename__ = "workflow_run_output_parameters"

-    workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), primary_key=True, index=True)
+    workflow_run_id = Column(
+        String,
+        ForeignKey("workflow_runs.workflow_run_id"),
+        primary_key=True,
+        index=True,
+    )
    output_parameter_id = Column(
-        String, ForeignKey("output_parameters.output_parameter_id"), primary_key=True, index=True
+        String,
+        ForeignKey("output_parameters.output_parameter_id"),
+        primary_key=True,
+        index=True,
    )
    value = Column(JSON, nullable=False)
    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
--- a/skyvern/forge/sdk/db/utils.py
+++ b/skyvern/forge/sdk/db/utils.py
@@ -67,7 +67,7 @@ def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False) -> Task:
        extracted_information=task_obj.extracted_information,
        failure_reason=task_obj.failure_reason,
        organization_id=task_obj.organization_id,
-        proxy_location=ProxyLocation(task_obj.proxy_location) if task_obj.proxy_location else None,
+        proxy_location=(ProxyLocation(task_obj.proxy_location) if task_obj.proxy_location else None),
        extracted_information_schema=task_obj.extracted_information_schema,
        workflow_run_id=task_obj.workflow_run_id,
        order=task_obj.order,
@@ -112,7 +112,9 @@ def convert_to_organization(org_model: OrganizationModel) -> Organization:
    )


-def convert_to_organization_auth_token(org_auth_token: OrganizationAuthTokenModel) -> OrganizationAuthToken:
+def convert_to_organization_auth_token(
+    org_auth_token: OrganizationAuthTokenModel,
+) -> OrganizationAuthToken:
    return OrganizationAuthToken(
        id=org_auth_token.id,
        organization_id=org_auth_token.organization_id,
@@ -126,7 +128,10 @@ def convert_to_organization_auth_token(org_auth_token: OrganizationAuthTokenMode

 def convert_to_artifact(artifact_model: ArtifactModel, debug_enabled: bool = False) -> Artifact:
    if debug_enabled:
-        LOG.debug("Converting ArtifactModel to Artifact", artifact_id=artifact_model.artifact_id)
+        LOG.debug(
+            "Converting ArtifactModel to Artifact",
+            artifact_id=artifact_model.artifact_id,
+        )

    return Artifact(
        artifact_id=artifact_model.artifact_id,
@@ -142,7 +147,10 @@ def convert_to_artifact(artifact_model: ArtifactModel, debug_enabled: bool = Fal

 def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = False) -> Workflow:
    if debug_enabled:
-        LOG.debug("Converting WorkflowModel to Workflow", workflow_id=workflow_model.workflow_id)
+        LOG.debug(
+            "Converting WorkflowModel to Workflow",
+            workflow_id=workflow_model.workflow_id,
+        )

    return Workflow(
        workflow_id=workflow_model.workflow_id,
@@ -150,7 +158,7 @@ def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = Fal
        title=workflow_model.title,
        workflow_permanent_id=workflow_model.workflow_permanent_id,
        webhook_callback_url=workflow_model.webhook_callback_url,
-        proxy_location=ProxyLocation(workflow_model.proxy_location) if workflow_model.proxy_location else None,
+        proxy_location=(ProxyLocation(workflow_model.proxy_location) if workflow_model.proxy_location else None),
        version=workflow_model.version,
        description=workflow_model.description,
        workflow_definition=WorkflowDefinition.model_validate(workflow_model.workflow_definition),
@@ -162,13 +170,18 @@ def convert_to_workflow(workflow_model: WorkflowModel, debug_enabled: bool = Fal

 def convert_to_workflow_run(workflow_run_model: WorkflowRunModel, debug_enabled: bool = False) -> WorkflowRun:
    if debug_enabled:
-        LOG.debug("Converting WorkflowRunModel to WorkflowRun", workflow_run_id=workflow_run_model.workflow_run_id)
+        LOG.debug(
+            "Converting WorkflowRunModel to WorkflowRun",
+            workflow_run_id=workflow_run_model.workflow_run_id,
+        )

    return WorkflowRun(
        workflow_run_id=workflow_run_model.workflow_run_id,
        workflow_id=workflow_run_model.workflow_id,
        status=WorkflowRunStatus[workflow_run_model.status],
-        proxy_location=ProxyLocation(workflow_run_model.proxy_location) if workflow_run_model.proxy_location else None,
+        proxy_location=(
+            ProxyLocation(workflow_run_model.proxy_location) if workflow_run_model.proxy_location else None
+        ),
        webhook_callback_url=workflow_run_model.webhook_callback_url,
        created_at=workflow_run_model.created_at,
        modified_at=workflow_run_model.modified_at,
@@ -221,7 +234,8 @@ def convert_to_aws_secret_parameter(


 def convert_to_bitwarden_login_credential_parameter(
-    bitwarden_login_credential_parameter_model: BitwardenLoginCredentialParameterModel, debug_enabled: bool = False
+    bitwarden_login_credential_parameter_model: BitwardenLoginCredentialParameterModel,
+    debug_enabled: bool = False,
 ) -> BitwardenLoginCredentialParameter:
    if debug_enabled:
        LOG.debug(