refactor: Extract thinking tokens and reasoning effort methods into separate functions

This commit is contained in:
Paul Gauthier (aider) 2025-03-11 12:09:41 -07:00
parent e21bab2d17
commit 79f714ab16
2 changed files with 38 additions and 26 deletions

View file

@ -209,34 +209,13 @@ class Coder:
output = f"{prefix}: {main_model.name} with {self.edit_format} edit format"
# Check for thinking token budget
if (
main_model.extra_params
and "thinking" in main_model.extra_params
and "budget_tokens" in main_model.extra_params["thinking"]
):
budget = main_model.extra_params["thinking"]["budget_tokens"]
# Format as xx.yK for thousands, xx.yM for millions
if budget >= 1024 * 1024:
value = budget / (1024 * 1024)
if value == int(value):
formatted_budget = f"{int(value)}M"
else:
formatted_budget = f"{value:.1f}M"
else:
value = budget / 1024
if value == int(value):
formatted_budget = f"{int(value)}k"
else:
formatted_budget = f"{value:.1f}k"
output += f", {formatted_budget} think tokens"
thinking_tokens = self.get_thinking_tokens(main_model)
if thinking_tokens:
output += f", {thinking_tokens} think tokens"
# Check for reasoning effort
if (
main_model.extra_params
and "extra_body" in main_model.extra_params
and "reasoning_effort" in main_model.extra_params["extra_body"]
):
reasoning_effort = main_model.extra_params["extra_body"]["reasoning_effort"]
reasoning_effort = self.get_reasoning_effort(main_model)
if reasoning_effort:
output += f", reasoning {reasoning_effort}"
if self.add_cache_headers or main_model.caches_by_default:

View file

@ -645,6 +645,39 @@ class Model(ModelSettings):
self.extra_params = {}
self.extra_params["thinking"] = {"type": "enabled", "budget_tokens": num_tokens}
def get_thinking_tokens(self, model):
"""Get formatted thinking token budget if available"""
if (
model.extra_params
and "thinking" in model.extra_params
and "budget_tokens" in model.extra_params["thinking"]
):
budget = model.extra_params["thinking"]["budget_tokens"]
# Format as xx.yK for thousands, xx.yM for millions
if budget >= 1024 * 1024:
value = budget / (1024 * 1024)
if value == int(value):
return f"{int(value)}M"
else:
return f"{value:.1f}M"
else:
value = budget / 1024
if value == int(value):
return f"{int(value)}k"
else:
return f"{value:.1f}k"
return None
def get_reasoning_effort(self, model):
"""Get reasoning effort value if available"""
if (
model.extra_params
and "extra_body" in model.extra_params
and "reasoning_effort" in model.extra_params["extra_body"]
):
return model.extra_params["extra_body"]["reasoning_effort"]
return None
def is_deepseek_r1(self):
name = self.name.lower()
if "deepseek" not in name: