Add task status as a query parameter (Skyvern-AI#229)

feiofil · Apr 24, 2024 · 5d90545 · 5d90545
1 parent 9d3c304
commit 5d90545
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py
@@ -390,7 +390,13 @@ async def update_task(
             LOG.error("UnexpectedError", exc_info=True)
             raise
 
-    async def get_tasks(self, page: int = 1, page_size: int = 10, organization_id: str | None = None) -> list[Task]:
+    async def get_tasks(
+        self,
+        page: int = 1,
+        page_size: int = 10,
+        task_status: TaskStatus | None = None,
+        organization_id: str | None = None,
+    ) -> list[Task]:
         """
         Get all tasks.
         :param page: Starts at 1
@@ -403,15 +409,11 @@ async def get_tasks(self, page: int = 1, page_size: int = 10, organization_id: s
         try:
             async with self.Session() as session:
                 db_page = page - 1  # offset logic is 0 based
-                tasks = (
-                    await session.scalars(
-                        select(TaskModel)
-                        .filter_by(organization_id=organization_id)
-                        .order_by(TaskModel.created_at.desc())
-                        .limit(page_size)
-                        .offset(db_page * page_size)
-                    )
-                ).all()
+                query = select(TaskModel).filter_by(organization_id=organization_id)
+                if task_status:
+                    query = query.filter_by(status=task_status)
+                query = query.order_by(TaskModel.created_at.desc()).limit(page_size).offset(db_page * page_size)
+                tasks = (await session.scalars(query)).all()
                 return [convert_to_task(task, debug_enabled=self.debug_enabled) for task in tasks]
         except SQLAlchemyError:
             LOG.error("SQLAlchemyError", exc_info=True)

diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -313,6 +313,7 @@ async def get_task_internal(
 async def get_agent_tasks(
     page: int = Query(1, ge=1),
     page_size: int = Query(10, ge=1),
+    task_status: TaskStatus | None = None,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> Response:
     """
@@ -323,7 +324,9 @@ async def get_agent_tasks(
         get_agent_task endpoint.
     """
     analytics.capture("skyvern-oss-agent-tasks-get")
-    tasks = await app.DATABASE.get_tasks(page, page_size, organization_id=current_org.organization_id)
+    tasks = await app.DATABASE.get_tasks(
+        page, page_size, task_status=task_status, organization_id=current_org.organization_id
+    )
     return ORJSONResponse([task.to_task_response().model_dump() for task in tasks])