Max retries per step configuration per org (#168)

skyvern/forge/sdk/db/client.py

@@ -418,12 +418,14 @@ class AgentDB:
         organization_name: str,
         webhook_callback_url: str | None = None,
         max_steps_per_run: int | None = None,
+        max_retries_per_step: int | None = None,
     ) -> Organization:
         async with self.Session() as session:
             org = OrganizationModel(
                 organization_name=organization_name,
                 webhook_callback_url=webhook_callback_url,
                 max_steps_per_run=max_steps_per_run,
+                max_retries_per_step=max_retries_per_step,
             )
             session.add(org)
             await session.commit()

skyvern/forge/sdk/db/models.py

@@ -74,7 +74,8 @@ class OrganizationModel(Base):
     organization_id = Column(String, primary_key=True, index=True, default=generate_org_id)
     organization_name = Column(String, nullable=False)
     webhook_callback_url = Column(UnicodeText)
-    max_steps_per_run = Column(Integer)
+    max_steps_per_run = Column(Integer, nullable=True)
+    max_retries_per_step = Column(Integer, nullable=True)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
     modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime, nullable=False)
 

skyvern/forge/sdk/db/utils.py

@@ -104,6 +104,7 @@ def convert_to_organization(org_model: OrganizationModel) -> Organization:
         organization_name=org_model.organization_name,
         webhook_callback_url=org_model.webhook_callback_url,
         max_steps_per_run=org_model.max_steps_per_run,
+        max_retries_per_step=org_model.max_retries_per_step,
         created_at=org_model.created_at,
         modified_at=org_model.modified_at,
     )

skyvern/forge/sdk/models.py

@@ -117,6 +117,7 @@ class Organization(BaseModel):
     organization_name: str
     webhook_callback_url: str | None = None
     max_steps_per_run: int | None = None
+    max_retries_per_step: int | None = None
 
     created_at: datetime
     modified_at: datetime