set extract-actions thinking budget to 512, others to 1024 (#4249)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user