From 965ff7c0b8f233ca681191acd21a42b9a1458525 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Fri, 23 Jan 2026 00:39:07 -0800 Subject: [PATCH] Make otp type also selectable in the totp endpoint (#4529) --- skyvern/forge/prompts/skyvern/parse-otp-login.j2 | 9 +++++++-- skyvern/forge/sdk/routes/credentials.py | 4 ++-- skyvern/forge/sdk/schemas/totp_codes.py | 15 ++++++++++----- skyvern/services/otp_service.py | 14 +++++++++++--- 4 files changed, 30 insertions(+), 12 deletions(-) diff --git a/skyvern/forge/prompts/skyvern/parse-otp-login.j2 b/skyvern/forge/prompts/skyvern/parse-otp-login.j2 index 033681a1..aafe81ff 100644 --- a/skyvern/forge/prompts/skyvern/parse-otp-login.j2 +++ b/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. diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py index bb1d0d69..784f47e4 100644 --- a/skyvern/forge/sdk/routes/credentials.py +++ b/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( diff --git a/skyvern/forge/sdk/schemas/totp_codes.py b/skyvern/forge/sdk/schemas/totp_codes.py index 238e4f45..0c5848ed 100644 --- a/skyvern/forge/sdk/schemas/totp_codes.py +++ b/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.") diff --git a/skyvern/services/otp_service.py b/skyvern/services/otp_service.py index 3041c531..fcd28399 100644 --- a/skyvern/services/otp_service.py +++ b/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)