From c939d34603e1a3860ee54e0fdc6378da1b7ed030 Mon Sep 17 00:00:00 2001
From: pedrohsdb <pedro@skyvern.com>
Date: Tue, 9 Dec 2025 12:57:10 -0800
Subject: [PATCH] set extract-actions thinking budget to 512, others to 1024
 (#4249)

---
 skyvern/config.py                                |  2 +-
 skyvern/forge/sdk/api/llm/api_handler_factory.py | 13 +++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/skyvern/config.py b/skyvern/config.py
index 0ea9d3b1..341c76c3 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -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
diff --git a/skyvern/forge/sdk/api/llm/api_handler_factory.py b/skyvern/forge/sdk/api/llm/api_handler_factory.py
index dd22f317..c432cb2d 100644
--- a/skyvern/forge/sdk/api/llm/api_handler_factory.py
+++ b/skyvern/forge/sdk/api/llm/api_handler_factory.py
@@ -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