Open source TOTP (#1677)

skyvern/forge/sdk/db/client.py

@@ -1924,6 +1924,33 @@ class AgentDB:
             totp_code = (await session.scalars(query)).all()
             return [TOTPCode.model_validate(totp_code) for totp_code in totp_code]
 
+    async def create_totp_code(
+        self,
+        organization_id: str,
+        totp_identifier: str,
+        content: str,
+        code: str,
+        task_id: str | None = None,
+        workflow_id: str | None = None,
+        source: str | None = None,
+        expired_at: datetime | None = None,
+    ) -> TOTPCode:
+        async with self.Session() as session:
+            new_totp_code = TOTPCodeModel(
+                organization_id=organization_id,
+                totp_identifier=totp_identifier,
+                content=content,
+                code=code,
+                task_id=task_id,
+                workflow_id=workflow_id,
+                source=source,
+                expired_at=expired_at,
+            )
+            session.add(new_totp_code)
+            await session.commit()
+            await session.refresh(new_totp_code)
+            return TOTPCode.model_validate(new_totp_code)
+
     async def create_action(self, action: Action) -> Action:
         async with self.Session() as session:
             new_action = ActionModel(

skyvern/forge/sdk/routes/totp.py

@@ -0,0 +1,45 @@
+import structlog
+from fastapi import APIRouter, Depends, HTTPException
+
+from skyvern.forge import app
+from skyvern.forge.prompts import prompt_engine
+from skyvern.forge.sdk.schemas.organizations import Organization
+from skyvern.forge.sdk.schemas.totp_codes import TOTPCode, TOTPCodeCreate
+from skyvern.forge.sdk.services import org_auth_service
+
+LOG = structlog.get_logger()
+totp_router = APIRouter()
+
+
+@totp_router.post("")
+@totp_router.post("/", include_in_schema=False)
+async def save_totp_code(
+    data: TOTPCodeCreate, curr_org: Organization = Depends(org_auth_service.get_current_org)
+) -> TOTPCode:
+    LOG.info(
+        "Saving TOTP code",
+        data=data,
+        organization_id=curr_org.organization_id,
+        totp_identifier=data.totp_identifier,
+        task_id=data.task_id,
+        workflow_id=data.workflow_id,
+    )
+    code = await parse_totp_code(data.content)
+    if not code:
+        raise HTTPException(status_code=400, detail="Failed to parse totp code")
+    return await app.DATABASE.create_totp_code(
+        organization_id=curr_org.organization_id,
+        totp_identifier=data.totp_identifier,
+        content=data.content,
+        code=code,
+        task_id=data.task_id,
+        workflow_id=data.workflow_id,
+        source=data.source,
+        expired_at=data.expired_at,
+    )
+
+
+async def parse_totp_code(content: str) -> str | None:
+    prompt = prompt_engine.load_prompt("parse-verification-code", content=content)
+    code_resp = await app.SECONDARY_LLM_API_HANDLER(prompt=prompt)
+    return code_resp.get("code", None)