From b4d5837196eb28e2553ad44bdc8e22cd8e09121c Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Sat, 31 May 2025 18:42:57 -0700
Subject: [PATCH] ModelSelector: add 'Skyvern Optimized' as hard choice (#2558)

---
 .../src/components/ModelSelector.tsx          | 18 ++++++++---
 skyvern/config.py                             | 24 +++++++-------
 skyvern/forge/agent.py                        |  3 +-
 skyvern/forge/sdk/schemas/task_v2.py          |  2 +-
 skyvern/forge/sdk/schemas/tasks.py            |  2 +-
 skyvern/forge/sdk/workflow/models/workflow.py | 31 ++++++++++++++++++-
 6 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/skyvern-frontend/src/components/ModelSelector.tsx b/skyvern-frontend/src/components/ModelSelector.tsx
index e4d2847d..d1774ec6 100644
--- a/skyvern-frontend/src/components/ModelSelector.tsx
+++ b/skyvern-frontend/src/components/ModelSelector.tsx
@@ -21,6 +21,10 @@ type Props = {
   onChange: (value: WorkflowModel | null) => void;
 };
 
+const constants = {
+  SkyvernOptimized: "Skyvern Optimized",
+} as const;
+
 function ModelSelector({
   clearable = true,
   value,
@@ -38,6 +42,7 @@ function ModelSelector({
   });
 
   const models = availableModels?.models ?? [];
+  const choices = [constants.SkyvernOptimized, ...models];
 
   return (
     <div className="flex items-center justify-between">
@@ -49,18 +54,23 @@ function ModelSelector({
         <Select
           value={value?.model ?? ""}
           onValueChange={(v) => {
-            onChange({ model: v });
+            const newValue = v === constants.SkyvernOptimized ? null : v;
+            onChange(newValue ? { model: newValue } : null);
           }}
         >
           <SelectTrigger
             className={(className || "") + (value && clearable ? " pr-10" : "")}
           >
-            <SelectValue placeholder="Skyvern Optimized" />
+            <SelectValue placeholder={constants.SkyvernOptimized} />
           </SelectTrigger>
           <SelectContent>
-            {models.map((m) => (
+            {choices.map((m) => (
               <SelectItem key={m} value={m}>
-                {m}
+                {m === constants.SkyvernOptimized ? (
+                  <span>Skyvern Optimized ✨</span>
+                ) : (
+                  m
+                )}
               </SelectItem>
             ))}
           </SelectContent>
diff --git a/skyvern/config.py b/skyvern/config.py
index 8a1b90b5..991d29f7 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -272,22 +272,22 @@ class Settings(BaseSettings):
 
         if self.is_cloud_environment():
             return {
-                "gemini-2.5-pro-preview-05-06": "VERTEX_GEMINI_2.5_PRO_PREVIEW",
-                "gemini-2.5-flash-preview-05-20": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
-                "azure/gpt-4.1": "AZURE_OPENAI_GPT4_1",
-                "azure/o3-mini": "AZURE_OPENAI_O3_MINI",
-                "us.anthropic.claude-opus-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_OPUS_INFERENCE_PROFILE",
-                "us.anthropic.claude-sonnet-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_SONNET_INFERENCE_PROFILE",
+                "Gemini 2.5": "GEMINI_2.5_PRO_PREVIEW",
+                "Gemini 2.5 Flash": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
+                "GPT 4.1": "OPENAI_GPT4_1",
+                "GPT o3-mini": "OPENAI_O3_MINI",
+                "bedrock/us.anthropic.claude-opus-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_OPUS_INFERENCE_PROFILE",
+                "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_SONNET_INFERENCE_PROFILE",
             }
         else:
             # TODO: apparently the list for OSS is to be much larger
             return {
-                "gemini-2.5-pro-preview-05-06": "VERTEX_GEMINI_2.5_PRO_PREVIEW",
-                "gemini-2.5-flash-preview-05-20": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
-                "azure/gpt-4.1": "AZURE_OPENAI_GPT4_1",
-                "azure/o3-mini": "AZURE_OPENAI_O3_MINI",
-                "us.anthropic.claude-opus-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_OPUS_INFERENCE_PROFILE",
-                "us.anthropic.claude-sonnet-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_SONNET_INFERENCE_PROFILE",
+                "Gemini 2.5": "GEMINI_2.5_PRO_PREVIEW",
+                "Gemini 2.5 Flash": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
+                "GPT 4.1": "OPENAI_GPT4_1",
+                "GPT o3-mini": "OPENAI_O3_MINI",
+                "bedrock/us.anthropic.claude-opus-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_OPUS_INFERENCE_PROFILE",
+                "bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0": "BEDROCK_ANTHROPIC_CLAUDE4_SONNET_INFERENCE_PROFILE",
             }
 
     def is_cloud_environment(self) -> bool:
diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 829fbc6c..a1d12632 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -867,13 +867,12 @@ class ForgeAgent:
                 else:
                     if engine in CUA_ENGINES:
                         self.async_operation_pool.run_operation(task.task_id, AgentPhase.llm)
-
                     json_response = await app.LLM_API_HANDLER(
                         prompt=extract_action_prompt,
                         prompt_name="extract-actions",
                         step=step,
                         screenshots=scraped_page.screenshots,
-                        llm_key_override=task.llm_key,
+                        llm_key_override=llm_caller.llm_key if llm_caller else None,
                     )
                     try:
                         json_response = await self.handle_potential_verification_code(
diff --git a/skyvern/forge/sdk/schemas/task_v2.py b/skyvern/forge/sdk/schemas/task_v2.py
index 57c8c516..9d76f001 100644
--- a/skyvern/forge/sdk/schemas/task_v2.py
+++ b/skyvern/forge/sdk/schemas/task_v2.py
@@ -57,7 +57,7 @@ class TaskV2(BaseModel):
         """
 
         if self.model:
-            model_name = self.model.get("model_name")
+            model_name = self.model.get("name")
             if model_name:
                 mapping = settings.get_model_name_to_llm_key()
                 llm_key = mapping.get(model_name)
diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py
index 47647984..c154012b 100644
--- a/skyvern/forge/sdk/schemas/tasks.py
+++ b/skyvern/forge/sdk/schemas/tasks.py
@@ -248,7 +248,7 @@ class Task(TaskBase):
         Otherwise return `None`.
         """
         if self.model:
-            model_name = self.model.get("model_name")
+            model_name = self.model.get("name")
             if model_name:
                 mapping = settings.get_model_name_to_llm_key()
                 return mapping.get(model_name)
diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py
index c3b8b7d0..254d6dac 100644
--- a/skyvern/forge/sdk/workflow/models/workflow.py
+++ b/skyvern/forge/sdk/workflow/models/workflow.py
@@ -5,10 +5,11 @@ from typing import Any, List
 from pydantic import BaseModel, field_validator
 from typing_extensions import deprecated
 
+from skyvern.config import settings
 from skyvern.forge.sdk.schemas.files import FileInfo
 from skyvern.forge.sdk.schemas.task_v2 import TaskV2
 from skyvern.forge.sdk.workflow.exceptions import WorkflowDefinitionHasDuplicateBlockLabels
-from skyvern.forge.sdk.workflow.models.block import BlockTypeVar
+from skyvern.forge.sdk.workflow.models.block import Block, BlockTypeVar
 from skyvern.forge.sdk.workflow.models.parameter import PARAMETER_TYPE
 from skyvern.schemas.runs import ProxyLocation
 from skyvern.utils.url_validators import validate_url
@@ -81,6 +82,34 @@ class Workflow(BaseModel):
     modified_at: datetime
     deleted_at: datetime | None = None
 
+    def determine_llm_key(self, *, block: Block | None = None) -> str | None:
+        """
+        Determine the LLM key override to use for a block, if it has one.
+
+        It has one if:
+            - it defines one, or
+            - the workflow it is a part of (if applicable) defines one
+        """
+
+        mapping = settings.get_model_name_to_llm_key()
+
+        if block:
+            model_name = (block.model or {}).get("name")
+
+            if model_name:
+                llm_key = mapping.get(model_name)
+                if llm_key:
+                    return llm_key
+
+        workflow_model_name = (self.model or {}).get("name")
+
+        if workflow_model_name:
+            llm_key = mapping.get(workflow_model_name)
+            if llm_key:
+                return llm_key
+
+        return None
+
 
 class WorkflowRunStatus(StrEnum):
     created = "created"