vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add task status as a query parameter (#229)

Browse Source
This commit is contained in:
Salih Altun
2024-04-24 20:39:19 +03:00
committed by GitHub
parent 9d3c3047c8
commit 5d9054594c
2 changed files with 16 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
skyvern/forge/sdk/db/client.py
View File

@@ -390,7 +390,13 @@ class AgentDB:
LOG.error("UnexpectedError", exc_info=True) LOG.error("UnexpectedError", exc_info=True)
raise 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. Get all tasks.
:param page: Starts at 1 :param page: Starts at 1
@@ -403,15 +409,11 @@ class AgentDB:
try: try:
async with self.Session() as session: async with self.Session() as session:
db_page = page - 1 # offset logic is 0 based db_page = page - 1 # offset logic is 0 based
tasks = ( query = select(TaskModel).filter_by(organization_id=organization_id)
await session.scalars( if task_status:
select(TaskModel) query = query.filter_by(status=task_status)
.filter_by(organization_id=organization_id) query = query.order_by(TaskModel.created_at.desc()).limit(page_size).offset(db_page * page_size)
.order_by(TaskModel.created_at.desc()) tasks = (await session.scalars(query)).all()
.limit(page_size)
.offset(db_page * page_size)
)
).all()
return [convert_to_task(task, debug_enabled=self.debug_enabled) for task in tasks] return [convert_to_task(task, debug_enabled=self.debug_enabled) for task in tasks]
except SQLAlchemyError: except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True) LOG.error("SQLAlchemyError", exc_info=True)

5
skyvern/forge/sdk/routes/agent_protocol.py
View File

@@ -313,6 +313,7 @@ async def get_task_internal(
async def get_agent_tasks( async def get_agent_tasks(
page: int = Query(1, ge=1), page: int = Query(1, ge=1),
page_size: int = Query(10, ge=1), page_size: int = Query(10, ge=1),
task_status: TaskStatus | None = None,
current_org: Organization = Depends(org_auth_service.get_current_org), current_org: Organization = Depends(org_auth_service.get_current_org),
) -> Response: ) -> Response:
""" """
@@ -323,7 +324,9 @@ async def get_agent_tasks(
get_agent_task endpoint. get_agent_task endpoint.
""" """
analytics.capture("skyvern-oss-agent-tasks-get") 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]) return ORJSONResponse([task.to_task_response().model_dump() for task in tasks])