From 800a26d323e3168ceb720075bba49018e5948b35 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Sat, 31 May 2025 18:40:36 -0700 Subject: [PATCH] use model_name instead of name (#2557) --- 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 +------------------ 5 files changed, 17 insertions(+), 45 deletions(-) diff --git a/skyvern/config.py b/skyvern/config.py index 991d29f7..8a1b90b5 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": "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": "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", } 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": "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", } 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/schemas/task_v2.py b/skyvern/forge/sdk/schemas/task_v2.py index 9d76f001..57c8c516 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("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) diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py index c154012b..47647984 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("name") + model_name = self.model.get("model_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 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"