Jon/model name massage (#2559)

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(

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(

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
 

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]

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"