Make otp type also selectable in the totp endpoint (#4529)

skyvern/forge/prompts/skyvern/parse-otp-login.j2

@@ -1,8 +1,13 @@
-You receive either an email or a text message containing an OTP(like TOTP, Magic Link) to verify the login. Your job is to parse the content, identify the OTP type and value. There should be only one OTP type and one OTP value in the content. The value must be from the content
+You receive either an email or a text message containing an OTP(like TOTP, Magic Link) to verify the login. Your job is to parse the content, identify the OTP type and value. The value must be from the content.
+{% if enforced_otp_type %}
+IMPORTANT: The user has specified they expect a "{{ enforced_otp_type }}" type OTP. You MUST extract the {{ enforced_otp_type }} value from the content, even if other OTP types are present. Ignore any other OTP types in the content. If no {{ enforced_otp_type }} is found, set otp_value_found to false.
+{% else %}
+There should be only one OTP type and one OTP value in the content.
+{% endif %}
 
 You should follow the rules below to identify the OTP type and value:
 - If it's a Magic Link login, the value is usually a link which must be a valid HTTP or HTTPS URL.
-- If it's a TOTP login, The most common value is a code which is a series of digits, although sometimes it may contain letters. 
+- If it's a TOTP login, The most common value is a code which is a series of digits, although sometimes it may contain letters.
 
 MAKE SURE YOU OUTPUT VALID JSON. No text before or after JSON, no trailing commas, no comments (//), no unnecessary quotes, etc.
 

skyvern/forge/sdk/routes/credentials.py

@@ -113,10 +113,10 @@ async def send_totp_code(
         if not workflow_run:
             raise HTTPException(status_code=400, detail=f"Invalid workflow run id: {data.workflow_run_id}")
     content = data.content.strip()
-    otp_value: OTPValue | None = OTPValue(value=content, type=OTPType.TOTP)
+    otp_value: OTPValue | None = OTPValue(value=content, type=data.type or OTPType.TOTP)
     # We assume the user is sending the code directly when the length of code is less than or equal to 10
     if len(content) > 10:
-        otp_value = await parse_otp_login(content, curr_org.organization_id)
+        otp_value = await parse_otp_login(content, curr_org.organization_id, enforced_otp_type=data.type)
 
     if not otp_value:
         LOG.error(

skyvern/forge/sdk/schemas/totp_codes.py

@@ -47,6 +47,11 @@ class TOTPCodeBase(BaseModel):
     )
 
 
+class OTPType(StrEnum):
+    TOTP = "totp"
+    MAGIC_LINK = "magic_link"
+
+
 class TOTPCodeCreate(TOTPCodeBase):
     totp_identifier: str = Field(
         ...,
@@ -58,6 +63,11 @@ class TOTPCodeCreate(TOTPCodeBase):
         description="The content of the TOTP code. It can be the email content that contains the TOTP code, or the sms message that contains the TOTP code. Skyvern will automatically extract the TOTP code from the content.",
         examples=["Hello, your verification code is 123456"],
     )
+    type: OTPType | None = Field(
+        default=None,
+        description="Optional. If provided, forces extraction of this specific OTP type (totp or magic_link). Use this when the content contains multiple OTP types and you want to specify which one to extract.",
+        examples=["totp", "magic_link"],
+    )
 
     @field_validator("content")
     @classmethod
@@ -66,11 +76,6 @@ class TOTPCodeCreate(TOTPCodeBase):
         return sanitize_postgres_text(value)
 
 
-class OTPType(StrEnum):
-    TOTP = "totp"
-    MAGIC_LINK = "magic_link"
-
-
 class TOTPCode(TOTPCodeCreate):
     totp_code_id: str = Field(..., description="The skyvern ID of the TOTP code.")
     code: str = Field(..., description="The TOTP code extracted from the content.")

skyvern/services/otp_service.py

@@ -36,12 +36,20 @@ class OTPResultParsedByLLM(BaseModel):
     otp_value: str | None = Field(None, description="The OTP value.")
 
 
-async def parse_otp_login(content: str, organization_id: str) -> OTPValue | None:
-    prompt = prompt_engine.load_prompt("parse-otp-login", content=content)
+async def parse_otp_login(
+    content: str,
+    organization_id: str,
+    enforced_otp_type: OTPType | None = None,
+) -> OTPValue | None:
+    prompt = prompt_engine.load_prompt(
+        "parse-otp-login",
+        content=content,
+        enforced_otp_type=enforced_otp_type.value if enforced_otp_type else None,
+    )
     resp = await app.SECONDARY_LLM_API_HANDLER(
         prompt=prompt, prompt_name="parse-otp-login", organization_id=organization_id
     )
-    LOG.info("OTP Login Parser Response", resp=resp)
+    LOG.info("OTP Login Parser Response", resp=resp, enforced_otp_type=enforced_otp_type)
     otp_result = OTPResultParsedByLLM.model_validate(resp)
     if otp_result.otp_value_found and otp_result.otp_value:
         return OTPValue(value=otp_result.otp_value, type=otp_result.otp_type)