set extract-actions thinking budget to 512, others to 1024 (#4249)

This commit is contained in:
pedrohsdb
2025-12-09 12:57:10 -08:00
committed by GitHub
parent 25c16edd85
commit c939d34603
2 changed files with 14 additions and 1 deletions

View File

@@ -291,7 +291,7 @@ class Settings(BaseSettings):
# GEMINI
GEMINI_API_KEY: str | None = None
GEMINI_INCLUDE_THOUGHT: bool = False
GEMINI_THINKING_BUDGET: int | None = 512
GEMINI_THINKING_BUDGET: int | None = None
# VERTEX_AI
VERTEX_CREDENTIALS: str | None = None

View File

@@ -49,6 +49,9 @@ LOG = structlog.get_logger()
EXTRACT_ACTION_PROMPT_NAME = "extract-actions"
CHECK_USER_GOAL_PROMPT_NAMES = {"check-user-goal", "check-user-goal-with-termination"}
# Default thinking budget for extract-actions prompt (can be overridden by THINKING_BUDGET_OPTIMIZATION experiment)
EXTRACT_ACTION_DEFAULT_THINKING_BUDGET = 512
@runtime_checkable
class RouterWithModelList(Protocol):
@@ -380,6 +383,11 @@ class LLMAPIHandlerFactory:
LLMAPIHandlerFactory._apply_thinking_budget_optimization(
parameters, new_budget, llm_config, prompt_name
)
elif prompt_name == EXTRACT_ACTION_PROMPT_NAME:
# Apply default thinking budget for extract-actions (512) unless overridden by experiment
LLMAPIHandlerFactory._apply_thinking_budget_optimization(
parameters, EXTRACT_ACTION_DEFAULT_THINKING_BUDGET, llm_config, prompt_name
)
context = skyvern_context.current()
is_speculative_step = step.is_speculative if step else False
@@ -780,6 +788,11 @@ class LLMAPIHandlerFactory:
LLMAPIHandlerFactory._apply_thinking_budget_optimization(
active_parameters, new_budget, llm_config, prompt_name
)
elif prompt_name == EXTRACT_ACTION_PROMPT_NAME:
# Apply default thinking budget for extract-actions (512) unless overridden by experiment
LLMAPIHandlerFactory._apply_thinking_budget_optimization(
active_parameters, EXTRACT_ACTION_DEFAULT_THINKING_BUDGET, llm_config, prompt_name
)
context = skyvern_context.current()
is_speculative_step = step.is_speculative if step else False