max task steps for task v2 (#1877)

This commit is contained in:
Shuchang Zheng
2025-03-04 01:07:07 -05:00
committed by GitHub
parent 618070840f
commit d31e4bf268
15 changed files with 90 additions and 40 deletions

View File

@@ -343,6 +343,26 @@ class AgentDB:
LOG.error("SQLAlchemyError", exc_info=True)
raise
async def get_total_step_count_by_task_ids(
self, task_ids: list[str], organization_id: str | None = None, statuses: list[StepStatus] | None = None
) -> int:
try:
async with self.Session() as session:
query = (
select(func.count())
.where(StepModel.task_id.in_(task_ids))
.filter_by(organization_id=organization_id)
)
if statuses:
query = query.filter(StepModel.status.in_(statuses))
return (await session.scalars(query)).scalar()
except SQLAlchemyError:
LOG.error("SQLAlchemyError", exc_info=True)
raise
except Exception:
LOG.error("UnexpectedError", exc_info=True)
raise
async def get_task_step_models(self, task_id: str, organization_id: str | None = None) -> Sequence[StepModel]:
try:
async with self.Session() as session: