Make rollout completion tokens configurable

skillopt/envs/spreadsheetbench/adapter.py

@@ -44,12 +44,15 @@ class SpreadsheetBenchAdapter(EnvAdapter):
         failure_only: bool = False,
         minibatch_size: int = 8,
         edit_budget: int = 4,
-        seed: int = 42,    ) -> None:
+        seed: int = 42,
+        max_completion_tokens: int = 16384,
+    ) -> None:
         self.data_root = data_root
         self.mode = mode  # "single", "multi", or "react"
         self.max_turns = max_turns
         self.exec_timeout = exec_timeout
         self.workers = workers
+        self.max_completion_tokens = int(max_completion_tokens)
         self.analyst_workers = analyst_workers
         self.failure_only = failure_only
         self.minibatch_size = minibatch_size
@@ -124,6 +127,7 @@ class SpreadsheetBenchAdapter(EnvAdapter):
                 skill_content=skill_content,
                 mode=self.mode,
                 max_turns=self.max_turns,
+                max_completion_tokens=self.max_completion_tokens,
                 max_api_workers=self.workers,
                 task_timeout=self.exec_timeout,
                 use_eval_feedback=kwargs.get("use_eval_feedback", False),
@@ -138,6 +142,7 @@ class SpreadsheetBenchAdapter(EnvAdapter):
                 out_root=out_dir,
                 skill_content=skill_content,
                 max_turns=self.max_turns,
+                max_completion_tokens=self.max_completion_tokens,
                 max_api_workers=self.workers,
                 task_timeout=max(600, int(self.exec_timeout) + 60),
                 diagnostic_mode=kwargs.get("diagnostic_mode", False),

skillopt/envs/spreadsheetbench/react_agent.py

@@ -365,7 +365,7 @@ def run_react(
     answer_position: str = "",
     skill_content: str = "",
     max_turns: int = 30,
-    max_output_tokens: int = 4096,
+    max_output_tokens: int = 16384,
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
     diagnostic_trace_context: str = "",

skillopt/envs/spreadsheetbench/rollout.py

@@ -174,6 +174,7 @@ def process_one(
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
     diagnostic_trace_context: str = "",
+    max_completion_tokens: int = 16384,
 ) -> dict:
     """Run the ReAct agent on a single SpreadsheetBench task.
 
@@ -283,6 +284,7 @@ def process_one(
                 answer_position=answer_position_eval,
                 skill_content=skill_content,
                 max_turns=max_turns,
+                max_output_tokens=max_completion_tokens,
                 diagnostic_mode=diagnostic_mode,
                 diagnostic_instruction=diagnostic_instruction,
                 diagnostic_trace_context=diagnostic_trace_context,
@@ -409,6 +411,7 @@ def run_spreadsheet_batch(
     out_root: str,
     skill_content: str,
     max_turns: int = 30,
+    max_completion_tokens: int = 16384,
     max_api_workers: int = 64,
     task_timeout: int = 600,
     diagnostic_mode: bool = False,
@@ -479,6 +482,7 @@ def run_spreadsheet_batch(
             diagnostic_mode,
             diagnostic_instruction,
             (diagnostic_trace_context_by_id or {}).get(str(it["id"]), ""),
+            max_completion_tokens,
         )
 
     ex = ThreadPoolExecutor(max_workers=max_api_workers)
@@ -542,6 +546,7 @@ def process_one_codegen(
     skill_content: str,
     mode: str = "single",
     max_turns: int = 5,
+    max_completion_tokens: int = 16384,
     use_eval_feedback: bool = False,
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
@@ -653,6 +658,7 @@ def process_one_codegen(
                     answer_position=answer_position_eval,
                     skill_content=skill_content,
                     max_turns=max_turns,
+                    max_output_tokens=max_completion_tokens,
                     gold_path=first_gold if use_eval_feedback else "",
                     diagnostic_mode=diagnostic_mode,
                     diagnostic_instruction=diagnostic_instruction,
@@ -666,6 +672,7 @@ def process_one_codegen(
                     instruction_type=instruction_type,
                     answer_position=answer_position_eval,
                     skill_content=skill_content,
+                    max_output_tokens=max_completion_tokens,
                     diagnostic_mode=diagnostic_mode,
                     diagnostic_instruction=diagnostic_instruction,
                     diagnostic_trace_context=diagnostic_trace_context,
@@ -790,6 +797,7 @@ def run_spreadsheet_batch_codegen(
     skill_content: str,
     mode: str = "single",
     max_turns: int = 5,
+    max_completion_tokens: int = 16384,
     max_api_workers: int = 32,
     task_timeout: int = 0,
     use_eval_feedback: bool = False,
@@ -845,6 +853,7 @@ def run_spreadsheet_batch_codegen(
             skill_content,
             mode,
             max_turns,
+            max_completion_tokens,
             use_eval_feedback,
             diagnostic_mode,
             diagnostic_instruction,