Switch to chat_template_kwargs for thinking

README.md

@@ -24,12 +24,11 @@ This project provides a small interface for running "tournaments" between langua
    - `ENABLE_GENERATE_THINKING`
    - `ENABLE_SCORE_THINKING`
    - `ENABLE_PAIRWISE_THINKING`
-   - `THINKING_BUDGET_TOKENS`
-   
+
    When any of the thinking flags are enabled, the app sends
-   `thinking={"type": "enabled", "budget_tokens": $THINKING_BUDGET_TOKENS}` with each
+   `chat_template_kwargs={"enable_thinking": True}` with each
    `litellm.completion` call for that model. Otherwise it sends
-   `thinking={"type": "disabled", "budget_tokens": 0}`.
+   `chat_template_kwargs={"enable_thinking": False}`.
 2. Install dependencies (example with `pip`):
    ```bash
    pip install gradio litellm python-dotenv tqdm matplotlib

main.py

@@ -48,7 +48,6 @@ PAIRWISE_WITH_INSTRUCTION_DEFAULT = os.getenv("PASS_INSTRUCTION_TO_PAIRWISE", "t
 GENERATE_THINKING_DEFAULT = os.getenv("ENABLE_GENERATE_THINKING", "false").lower() == "true"
 SCORE_THINKING_DEFAULT = os.getenv("ENABLE_SCORE_THINKING", "false").lower() == "true"
 PAIRWISE_THINKING_DEFAULT = os.getenv("ENABLE_PAIRWISE_THINKING", "false").lower() == "true"
-THINKING_BUDGET_TOKENS_DEFAULT = int(os.getenv("THINKING_BUDGET_TOKENS", "1024"))
 CRITERIA_DEFAULT = "Factuality,Instruction Following,Precision"
 def _clean_json(txt):
     txt = re.sub(r"^```.*?\n|```$", "", txt, flags=re.DOTALL).strip()
@@ -164,7 +163,6 @@ def run_tournament(
         api_key=api_token,
         temperature=generate_temperature,
         thinking=generate_thinking,
-        budget_tokens=THINKING_BUDGET_TOKENS_DEFAULT,
         return_usage=True,
     )
     add_usage(usage)
@@ -190,7 +188,6 @@ def run_tournament(
                 temperature=score_temperature,
                 include_instruction=score_with_instruction,
                 thinking=score_thinking,
-                budget_tokens=THINKING_BUDGET_TOKENS_DEFAULT,
                 return_usage=True,
             )
             add_usage(usage)
@@ -230,7 +227,6 @@ def run_tournament(
                 temperature=pairwise_temperature,
                 include_instruction=pairwise_with_instruction,
                 thinking=pairwise_thinking,
-                budget_tokens=THINKING_BUDGET_TOKENS_DEFAULT,
                 return_usage=True,
             )
             add_usage(usage)

tests/test_main.py

@@ -115,7 +115,7 @@ def test_run_tournament_full_loop():
     assert 'Done' in process_log
     assert hist_fig == 'fig'
     assert top_picks.strip() in {'p1', 'p2'}
-    mock_gen.assert_called_once_with('instr', 4, model='gm', api_base='b', api_key='k', temperature=1, thinking=True, budget_tokens=1024, return_usage=True)
+    mock_gen.assert_called_once_with('instr', 4, model='gm', api_base='b', api_key='k', temperature=1, thinking=True, return_usage=True)
     assert 'Score completion' in process_log
     assert 'Pairwise completion' in process_log
     assert 'Prompt tokens' in usage

tests/test_tournament_utils.py

@@ -26,7 +26,7 @@ def test_generate_players():
     resp = make_response([" player1 ", "player2\n"])
     with patch('tournament_utils.completion', return_value=resp) as mock_comp:
         players = tu.generate_players('instr', 2, model='m', api_base='b', api_key='k', temperature=0.5)
-        mock_comp.assert_called_once_with(model='m', messages=[{'role': 'user', 'content': 'instr'}], n=2, api_base='b', api_key='k', temperature=0.5, thinking={'type': 'disabled', 'budget_tokens': 0})
+        mock_comp.assert_called_once_with(model='m', messages=[{'role': 'user', 'content': 'instr'}], n=2, api_base='b', api_key='k', temperature=0.5, chat_template_kwargs={'enable_thinking': False})
         assert players == ['player1', 'player2']
 
 
@@ -60,7 +60,7 @@ def test_thinking_passed_to_completion():
         tu.prompt_pairwise('i', 'block', 'a', 'b', thinking=True)
         assert mock_comp.call_count == 3
         for call in mock_comp.call_args_list:
-            assert call.kwargs['thinking'] == {'type': 'enabled', 'budget_tokens': 1024}
+            assert call.kwargs['chat_template_kwargs'] == {'enable_thinking': True}
 
 
 def test_thinking_disabled_by_default():
@@ -71,4 +71,4 @@ def test_thinking_disabled_by_default():
         tu.prompt_pairwise('i', 'block', 'a', 'b')
         assert mock_comp.call_count == 3
         for call in mock_comp.call_args_list:
-            assert call.kwargs['thinking'] == {'type': 'disabled', 'budget_tokens': 0}
+            assert call.kwargs['chat_template_kwargs'] == {'enable_thinking': False}

tournament_utils.py

@@ -1,8 +1,5 @@
-import os
 from litellm import completion
 
-BUDGET_TOKENS_DEFAULT = int(os.getenv("THINKING_BUDGET_TOKENS", "1024"))
-
 
 def _completion_kwargs(
     api_base: str | None,
@@ -29,7 +26,6 @@ def generate_players(
     api_key: str | None = None,
     temperature: float | None = None,
     thinking: bool = False,
-    budget_tokens: int = BUDGET_TOKENS_DEFAULT,
     return_usage: bool = False,
 ) -> list[str] | tuple[list[str], object]:
     """Request ``n`` completions for the instruction using the given model.
@@ -39,10 +35,7 @@ def generate_players(
     """
     messages = [{"role": "user", "content": instruction}]
     kwargs = _completion_kwargs(api_base, api_key, temperature)
-    kwargs["thinking"] = {
-        "type": "enabled" if thinking else "disabled",
-        "budget_tokens": budget_tokens if thinking else 0,
-    }
+    kwargs["chat_template_kwargs"] = {"enable_thinking": thinking}
     response = completion(
         model=model,
         messages=messages,
@@ -67,7 +60,6 @@ def prompt_score(
     temperature: float | None = None,
     include_instruction: bool = True,
     thinking: bool = False,
-    budget_tokens: int = BUDGET_TOKENS_DEFAULT,
     return_usage: bool = False,
 ) -> str | tuple[str, object]:
     """Return a JSON score string evaluating `player` on the criteria."""
@@ -80,10 +72,7 @@ Return JSON exactly like: {{"scores": [{example_scores}]}}."""
         prompt += f"\n\nInstruction:\n{instruction}"
     prompt += f"\n\nOutput:\n{player}"
     kwargs = _completion_kwargs(api_base, api_key, temperature)
-    kwargs["thinking"] = {
-        "type": "enabled" if thinking else "disabled",
-        "budget_tokens": budget_tokens if thinking else 0,
-    }
+    kwargs["chat_template_kwargs"] = {"enable_thinking": thinking}
     response = completion(
         model=model,
         messages=[{"role": "system", "content": prompt}],
@@ -107,7 +96,6 @@ def prompt_pairwise(
     temperature: float | None = None,
     include_instruction: bool = True,
     thinking: bool = False,
-    budget_tokens: int = BUDGET_TOKENS_DEFAULT,
     return_usage: bool = False,
 ) -> str | tuple[str, object]:
     """Return which player wins in JSON using the given criteria."""
@@ -119,10 +107,7 @@ Return ONLY JSON {{"winner": "A"}} or {{"winner": "B"}}."""
         prompt += f"\n\nInstruction:\n{instruction}"
     prompt += f"\n\nPlayers:\n<A>{a}</A>\n<B>{b}</B>"
     kwargs = _completion_kwargs(api_base, api_key, temperature)
-    kwargs["thinking"] = {
-        "type": "enabled" if thinking else "disabled",
-        "budget_tokens": budget_tokens if thinking else 0,
-    }
+    kwargs["chat_template_kwargs"] = {"enable_thinking": thinking}
     response = completion(
         model=model,
         messages=[{"role": "system", "content": prompt}],