diff --git a/skyvern/forge/sdk/db/agent_db.py b/skyvern/forge/sdk/db/agent_db.py index 34a417f8..a1660d4f 100644 --- a/skyvern/forge/sdk/db/agent_db.py +++ b/skyvern/forge/sdk/db/agent_db.py @@ -560,6 +560,26 @@ class AgentDB(BaseAlchemyDB): LOG.error("UnexpectedError", exc_info=True) raise + async def get_action_count_for_step(self, step_id: str, task_id: str, organization_id: str) -> int: + """Get count of actions for a step. Uses composite index for efficiency.""" + try: + async with self.Session() as session: + query = ( + select(func.count()) + .select_from(ActionModel) + .where(ActionModel.organization_id == organization_id) + .where(ActionModel.task_id == task_id) + .where(ActionModel.step_id == step_id) + ) + result = await session.scalar(query) + return result or 0 + except SQLAlchemyError: + LOG.error("SQLAlchemyError", exc_info=True) + raise + except Exception: + LOG.error("UnexpectedError", exc_info=True) + raise + async def get_first_step(self, task_id: str, organization_id: str | None = None) -> Step | None: try: async with self.Session() as session: diff --git a/skyvern/forge/sdk/db/id.py b/skyvern/forge/sdk/db/id.py index 1af0bafd..98c8f917 100644 --- a/skyvern/forge/sdk/db/id.py +++ b/skyvern/forge/sdk/db/id.py @@ -68,6 +68,7 @@ WORKFLOW_RUN_BLOCK_PREFIX = "wrb" WORKFLOW_RUN_PREFIX = "wr" WORKFLOW_SCRIPT_PREFIX = "ws" WORKFLOW_TEMPLATE_PREFIX = "wt" +ORGANIZATION_BILLING_PREFIX = "ob" def generate_workflow_id() -> str: @@ -260,6 +261,11 @@ def generate_script_block_id() -> str: return f"{SCRIPT_BLOCK_PREFIX}_{int_id}" +def generate_billing_id() -> str: + int_id = generate_id() + return f"{ORGANIZATION_BILLING_PREFIX}_{int_id}" + + ############# Helper functions below ############## def generate_id() -> int: """