diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index f785000b..eec6a97d 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -765,10 +765,13 @@ async def user_code(): return self.build_block_result(success=True, output_parameter_value=result, status=BlockStatus.completed) +DEFAULT_TEXT_PROMPT_LLM_KEY = settings.SECONDARY_LLM_KEY or settings.LLM_KEY + + class TextPromptBlock(Block): block_type: Literal[BlockType.TEXT_PROMPT] = BlockType.TEXT_PROMPT - llm_key: str + llm_key: str = DEFAULT_TEXT_PROMPT_LLM_KEY prompt: str parameters: list[PARAMETER_TYPE] = [] json_schema: dict[str, Any] | None = None diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py index c5c048ef..8c6b847e 100644 --- a/skyvern/forge/sdk/workflow/models/yaml.py +++ b/skyvern/forge/sdk/workflow/models/yaml.py @@ -3,6 +3,7 @@ from typing import Annotated, Any, Literal from pydantic import BaseModel, Field +from skyvern.config import settings from skyvern.forge.sdk.schemas.tasks import ProxyLocation from skyvern.forge.sdk.workflow.models.block import BlockType, FileType from skyvern.forge.sdk.workflow.models.parameter import ParameterType, WorkflowParameterType @@ -154,6 +155,9 @@ class CodeBlockYAML(BlockYAML): parameter_keys: list[str] | None = None +DEFAULT_TEXT_PROMPT_LLM_KEY = settings.SECONDARY_LLM_KEY or settings.LLM_KEY + + class TextPromptBlockYAML(BlockYAML): # There is a mypy bug with Literal. Without the type: ignore, mypy will raise an error: # Parameter 1 of Literal[...] cannot be of type "Any" @@ -161,7 +165,7 @@ class TextPromptBlockYAML(BlockYAML): # to infer the type of the parameter_type attribute. block_type: Literal[BlockType.TEXT_PROMPT] = BlockType.TEXT_PROMPT # type: ignore - llm_key: str + llm_key: str = DEFAULT_TEXT_PROMPT_LLM_KEY prompt: str parameter_keys: list[str] | None = None json_schema: dict[str, Any] | None = None