From 2167d88c203760c64cd3172e868304140f498d57 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Sat, 31 May 2025 19:34:30 -0700
Subject: [PATCH] Jon/model name massage (#2559)

---
 skyvern-frontend/src/api/types.ts             |  2 +-
 .../src/components/ModelSelector.tsx          | 20 +++++++--
 .../editor/nodes/NavigationNode/types.ts      |  2 +-
 .../editor/nodes/StartNode/StartNode.tsx      |  9 +++-
 .../routes/workflows/types/workflowTypes.ts   |  2 +-
 skyvern/config.py                             | 44 +++++++++++++------
 skyvern/forge/agent.py                        |  3 +-
 skyvern/forge/sdk/routes/agent_protocol.py    |  4 +-
 skyvern/forge/sdk/schemas/task_v2.py          |  4 +-
 skyvern/forge/sdk/schemas/tasks.py            |  6 +--
 skyvern/forge/sdk/workflow/models/workflow.py | 31 +------------
 11 files changed, 67 insertions(+), 60 deletions(-)

diff --git a/skyvern-frontend/src/api/types.ts b/skyvern-frontend/src/api/types.ts
index b54ebc3e..cef3eb67 100644
--- a/skyvern-frontend/src/api/types.ts
+++ b/skyvern-frontend/src/api/types.ts
@@ -393,7 +393,7 @@ export type CreditCardCredential = {
 };
 
 export type ModelsResponse = {
-  models: string[];
+  models: Record<string, string>;
 };
 
 export const RunEngine = {
diff --git a/skyvern-frontend/src/components/ModelSelector.tsx b/skyvern-frontend/src/components/ModelSelector.tsx
index d1774ec6..69185586 100644
--- a/skyvern-frontend/src/components/ModelSelector.tsx
+++ b/skyvern-frontend/src/components/ModelSelector.tsx
@@ -41,8 +41,17 @@ function ModelSelector({
     },
   });
 
-  const models = availableModels?.models ?? [];
-  const choices = [constants.SkyvernOptimized, ...models];
+  const models = availableModels?.models ?? {};
+  const reverseMap = Object.entries(models).reduce(
+    (acc, [key, value]) => {
+      acc[value] = key;
+      return acc;
+    },
+    {} as Record<string, string>,
+  );
+  const labels = Object.keys(reverseMap);
+  const chosen = value ? models[value.model_name] : constants.SkyvernOptimized;
+  const choices = [constants.SkyvernOptimized, ...labels];
 
   return (
     <div className="flex items-center justify-between">
@@ -52,10 +61,13 @@ function ModelSelector({
       </div>
       <div className="relative flex items-center">
         <Select
-          value={value?.model ?? ""}
+          value={chosen}
           onValueChange={(v) => {
             const newValue = v === constants.SkyvernOptimized ? null : v;
-            onChange(newValue ? { model: newValue } : null);
+            const modelName = newValue ? reverseMap[newValue] : null;
+            const value = modelName ? { model_name: modelName } : null;
+            console.log({ v, newValue, modelName, value });
+            onChange(value);
           }}
         >
           <SelectTrigger
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/NavigationNode/types.ts b/skyvern-frontend/src/routes/workflows/editor/nodes/NavigationNode/types.ts
index 3457aaee..8c6e2291 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/NavigationNode/types.ts
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/NavigationNode/types.ts
@@ -29,7 +29,7 @@ export const navigationNodeDefaultData: NavigationNodeData = {
   completeCriterion: "",
   terminateCriterion: "",
   errorCodeMapping: "null",
-  model: { model: "" },
+  model: { model_name: "" },
   engine: RunEngine.SkyvernV1,
   maxRetries: null,
   maxStepsOverride: null,
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx
index e1fbfac4..b2a2774d 100644
--- a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx
@@ -19,6 +19,7 @@ import { Switch } from "@/components/ui/switch";
 import { Separator } from "@/components/ui/separator";
 import { ModelsResponse } from "@/api/types";
 import { ModelSelector } from "@/components/ModelSelector";
+import { WorkflowModel } from "@/routes/workflows/types/workflowTypes";
 
 function StartNode({ id, data }: NodeProps<StartNode>) {
   const credentialGetter = useCredentialGetter();
@@ -33,7 +34,11 @@ function StartNode({ id, data }: NodeProps<StartNode>) {
     },
   });
 
-  const models = availableModels?.models ?? [];
+  const modelNames = availableModels?.models ?? {};
+  const firstKey = Object.keys(modelNames)[0];
+  const workflowModel: WorkflowModel | null = firstKey
+    ? { model_name: modelNames[firstKey] || "" }
+    : null;
 
   const [inputs, setInputs] = useState({
     webhookCallbackUrl: data.withWorkflowSettings
@@ -45,7 +50,7 @@ function StartNode({ id, data }: NodeProps<StartNode>) {
     persistBrowserSession: data.withWorkflowSettings
       ? data.persistBrowserSession
       : false,
-    model: data.withWorkflowSettings ? data.model : { model: models[0] || "" },
+    model: data.withWorkflowSettings ? data.model : workflowModel,
   });
 
   function handleChange(key: string, value: unknown) {
diff --git a/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts b/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts
index dd656d37..00949679 100644
--- a/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts
+++ b/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts
@@ -466,7 +466,7 @@ export type WorkflowSettings = {
   model: WorkflowModel | null;
 };
 
-export type WorkflowModel = JsonObjectExtendable<{ model: string }>;
+export type WorkflowModel = JsonObjectExtendable<{ model_name: string }>;
 
 export function isOutputParameter(
   parameter: Parameter,
diff --git a/skyvern/config.py b/skyvern/config.py
index 991d29f7..4b38b570 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -264,7 +264,7 @@ class Settings(BaseSettings):
     SKYVERN_BASE_URL: str = "https://api.skyvern.com"
     SKYVERN_API_KEY: str = "PLACEHOLDER"
 
-    def get_model_name_to_llm_key(self) -> dict[str, str]:
+    def get_model_name_to_llm_key(self) -> dict[str, dict[str, str]]:
         """
         Keys are model names available to blocks in the frontend. These map to key names
         in LLMConfigRegistry._configs.
@@ -272,22 +272,40 @@ class Settings(BaseSettings):
 
         if self.is_cloud_environment():
             return {
-                "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",
+                "gemini-2.5-pro-preview-05-06": {"llm_key": "VERTEX_GEMINI_2.5_PRO_PREVIEW", "label": "Gemini 2.5 Pro"},
+                "gemini-2.5-flash-preview-05-20": {
+                    "llm_key": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
+                    "label": "Gemini 2.5 Flash",
+                },
+                "azure/gpt-4.1": {"llm_key": "AZURE_OPENAI_GPT4_1", "label": "GPT 4.1"},
+                "azure/o3-mini": {"llm_key": "AZURE_OPENAI_O3_MINI", "label": "GPT O3 Mini"},
+                "us.anthropic.claude-opus-4-20250514-v1:0": {
+                    "llm_key": "BEDROCK_ANTHROPIC_CLAUDE4_OPUS_INFERENCE_PROFILE",
+                    "label": "Anthropic Claude 4 Opus",
+                },
+                "us.anthropic.claude-sonnet-4-20250514-v1:0": {
+                    "llm_key": "BEDROCK_ANTHROPIC_CLAUDE4_SONNET_INFERENCE_PROFILE",
+                    "label": "Anthropic Claude 4 Sonnet",
+                },
             }
         else:
             # TODO: apparently the list for OSS is to be much larger
             return {
-                "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",
+                "gemini-2.5-pro-preview-05-06": {"llm_key": "VERTEX_GEMINI_2.5_PRO_PREVIEW", "label": "Gemini 2.5 Pro"},
+                "gemini-2.5-flash-preview-05-20": {
+                    "llm_key": "VERTEX_GEMINI_2.5_FLASH_PREVIEW_05_20",
+                    "label": "Gemini 2.5 Flash",
+                },
+                "azure/gpt-4.1": {"llm_key": "AZURE_OPENAI_GPT4_1", "label": "GPT 4.1"},
+                "azure/o3-mini": {"llm_key": "AZURE_OPENAI_O3_MINI", "label": "GPT O3 Mini"},
+                "us.anthropic.claude-opus-4-20250514-v1:0": {
+                    "llm_key": "BEDROCK_ANTHROPIC_CLAUDE4_OPUS_INFERENCE_PROFILE",
+                    "label": "Anthropic Claude 4 Opus",
+                },
+                "us.anthropic.claude-sonnet-4-20250514-v1:0": {
+                    "llm_key": "BEDROCK_ANTHROPIC_CLAUDE4_SONNET_INFERENCE_PROFILE",
+                    "label": "Anthropic Claude 4 Sonnet",
+                },
             }
 
     def is_cloud_environment(self) -> bool:
diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index a1d12632..829fbc6c 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -867,12 +867,13 @@ 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=llm_caller.llm_key if llm_caller else None,
+                        llm_key_override=task.llm_key,
                     )
                     try:
                         json_response = await self.handle_potential_verification_code(
diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py
index ff5e9826..367c94d4 100644
--- a/skyvern/forge/sdk/routes/agent_protocol.py
+++ b/skyvern/forge/sdk/routes/agent_protocol.py
@@ -810,9 +810,9 @@ async def models() -> ModelsResponse:
     Get a list of available models.
     """
     mapping = settings.get_model_name_to_llm_key()
-    models = list(mapping.keys())
+    just_labels = {k: v["label"] for k, v in mapping.items() if "anthropic" not in k.lower()}
 
-    return ModelsResponse(models=models)
+    return ModelsResponse(models=just_labels)
 
 
 @legacy_base_router.post(
diff --git a/skyvern/forge/sdk/schemas/task_v2.py b/skyvern/forge/sdk/schemas/task_v2.py
index 9d76f001..1bf1e26b 100644
--- a/skyvern/forge/sdk/schemas/task_v2.py
+++ b/skyvern/forge/sdk/schemas/task_v2.py
@@ -57,10 +57,10 @@ class TaskV2(BaseModel):
         """
 
         if self.model:
-            model_name = self.model.get("name")
+            model_name = self.model.get("model_name")
             if model_name:
                 mapping = settings.get_model_name_to_llm_key()
-                llm_key = mapping.get(model_name)
+                llm_key = mapping.get(model_name, {}).get("llm_key")
                 if llm_key:
                     return llm_key
 
diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py
index c154012b..e09f5b0c 100644
--- a/skyvern/forge/sdk/schemas/tasks.py
+++ b/skyvern/forge/sdk/schemas/tasks.py
@@ -248,10 +248,10 @@ class Task(TaskBase):
         Otherwise return `None`.
         """
         if self.model:
-            model_name = self.model.get("name")
+            model_name = self.model.get("model_name")
             if model_name:
                 mapping = settings.get_model_name_to_llm_key()
-                return mapping.get(model_name)
+                return mapping.get(model_name, {}).get("llm_key")
 
         return None
 
@@ -370,4 +370,4 @@ class SortDirection(StrEnum):
 
 
 class ModelsResponse(BaseModel):
-    models: list[str]
+    models: dict[str, str]
diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py
index 254d6dac..c3b8b7d0 100644
--- a/skyvern/forge/sdk/workflow/models/workflow.py
+++ b/skyvern/forge/sdk/workflow/models/workflow.py
@@ -5,11 +5,10 @@ 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 Block, BlockTypeVar
+from skyvern.forge.sdk.workflow.models.block import 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
@@ -82,34 +81,6 @@ 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"