Make rollout completion tokens configurable

skillopt/envs/searchqa/adapter.py

@@ -31,10 +31,13 @@ class SearchQAAdapter(EnvAdapter):
         minibatch_size: int = 8,
         edit_budget: int = 4,
         seed: int = 42,
-        limit: int = 0,    ) -> None:
+        limit: int = 0,
+        max_completion_tokens: int = 16384,
+    ) -> None:
         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
@@ -84,6 +87,7 @@ class SearchQAAdapter(EnvAdapter):
             max_turns=self.max_turns,
             exec_timeout=self.exec_timeout,
             workers=self.workers,
+            max_completion_tokens=self.max_completion_tokens,
             diagnostic_mode=kwargs.get("diagnostic_mode", False),
             diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
             diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),

skillopt/envs/searchqa/rollout.py

@@ -148,6 +148,7 @@ def process_one(
     diagnostic_instruction: str = "",
     diagnostic_trace_context: str = "",
     exec_timeout: int = 120,
+    max_completion_tokens: int = 16384,
 ) -> dict:
     """Process a single QA item: run agent + evaluate.
 
@@ -268,7 +269,7 @@ def process_one(
             if turn == 0:
                 resp_text, _ = chat_target(
                     system=system, user=user,
-                    max_completion_tokens=512,
+                    max_completion_tokens=max_completion_tokens,
                     retries=5, stage="rollout",
                     timeout=exec_timeout,
                 )
@@ -281,7 +282,7 @@ def process_one(
                 )
                 resp_text, _ = chat_target(
                     system=system, user=refinement,
-                    max_completion_tokens=512,
+                    max_completion_tokens=max_completion_tokens,
                     retries=5, stage="rollout",
                     timeout=exec_timeout,
                 )
@@ -352,6 +353,7 @@ def run_batch(
     max_turns: int = 1,
     exec_timeout: int = 120,
     workers: int = 64,
+    max_completion_tokens: int = 16384,
     diagnostic_mode: bool = False,
     diagnostic_instruction: str = "",
     diagnostic_trace_context_by_id: dict[str, str] | None = None,
@@ -423,6 +425,7 @@ def run_batch(
             diagnostic_instruction,
             (diagnostic_trace_context_by_id or {}).get(str(item["id"]), ""),
             exec_timeout,
+            max_completion_tokens,
         )
 
     with open(results_path, "a") as outf: