Enhance LLM configuration with thinking budget support

- Added optional `thinking_budget` parameter to `config.yaml` for Gemini model, allowing for reserved thinking tokens.
- Updated `LlmSettings` class in `config.py` to include `thinking_budget` and adjusted parsing logic accordingly.
- Modified LLM call logic in `llm.py` to utilize `thinking_budget` for calculating `max_output_tokens`, improving output management.
This commit is contained in:
Arsham Mirehvandi 2026-09-17 09:44:04 +02:00
parent 253a6d853d
commit 958a17bb1f
3 changed files with 25 additions and 4 deletions

View File

@ -2,6 +2,9 @@
# Two independent profiles, one per pipeline stage. Each accepts the same keys:
# provider: openai | anthropic | gemini
# openai_model / anthropic_model / gemini_model / max_tokens / temperature
# thinking_budget (optional, Gemini only): reserved thinking tokens. When set,
# Gemini's max_output_tokens is max_tokens + thinking_budget, so max_tokens
# is the visible-output budget. Omit it to keep max_tokens as the shared cap.
llm:
# Part_One: vector-search query synthesis (small structured JSON in,
# 1-2 short queries out). A lighter/cheaper model is sufficient here.
@ -22,7 +25,8 @@ llm:
openai_model: gpt-4o
anthropic_model: claude-opus-4-5
gemini_model: gemini-2.5-pro
max_tokens: 16384
max_tokens: 4096
thinking_budget: 9216
temperature: 0.0
# ── Crop-first worklist ────────────────────────────────────────────────────────

View File

@ -20,6 +20,7 @@ class LlmSettings:
gemini_model: str
max_tokens: int
temperature: float
thinking_budget: int | None
@dataclass(frozen=True)
@ -119,6 +120,7 @@ class Settings:
def _parse_llm_settings(raw: dict, defaults: dict) -> LlmSettings:
"""Build an LlmSettings profile from a config sub-block, falling back to defaults."""
thinking_raw = raw["thinking_budget"] if "thinking_budget" in raw else defaults.get("thinking_budget")
return LlmSettings(
provider=str(raw.get("provider", defaults["provider"])).strip().lower(),
openai_model=str(raw.get("openai_model", defaults["openai_model"])),
@ -126,6 +128,7 @@ def _parse_llm_settings(raw: dict, defaults: dict) -> LlmSettings:
gemini_model=str(raw.get("gemini_model", defaults["gemini_model"])),
max_tokens=int(raw.get("max_tokens", defaults["max_tokens"])),
temperature=float(raw.get("temperature", defaults["temperature"])),
thinking_budget=None if thinking_raw is None else int(thinking_raw),
)
@ -235,14 +238,16 @@ def load_settings(
"gemini_model": "gemini-2.5-flash",
"max_tokens": 8192,
"temperature": 0.0,
"thinking_budget": None,
}
advice_generation_defaults = {
"provider": "gemini",
"openai_model": "gpt-4o",
"anthropic_model": "claude-opus-4-5",
"gemini_model": "gemini-2.5-pro",
"max_tokens": 65536,
"max_tokens": 8192,
"temperature": 0.0,
"thinking_budget": 8192,
}
return Settings(

View File

@ -101,7 +101,12 @@ def call_llm(
"""Dispatch a system/user prompt pair to the provider configured in llm_cfg."""
provider = llm_cfg.provider
logger.info("Calling LLM provider=%s model=%s", provider, _model_name(llm_cfg))
logger.info(
"Calling LLM provider=%s model=%s thinking_budget=%s",
provider,
_model_name(llm_cfg),
llm_cfg.thinking_budget,
)
logger.debug(
"LLM REQUEST (system_prompt + JSON user content):\n--- SYSTEM PROMPT ---\n%s\n--- USER JSON ---\n%s",
system_prompt,
@ -140,6 +145,12 @@ def _call_gemini(
from google import genai
from google.genai import types
max_output_tokens = llm_cfg.max_tokens
thinking_config = None
if llm_cfg.thinking_budget is not None:
max_output_tokens = llm_cfg.max_tokens + llm_cfg.thinking_budget
thinking_config = types.ThinkingConfig(thinking_budget=llm_cfg.thinking_budget)
client = genai.Client(api_key=settings.gemini_api_key)
response = client.models.generate_content(
model=llm_cfg.gemini_model,
@ -147,7 +158,8 @@ def _call_gemini(
config=types.GenerateContentConfig(
system_instruction=system_prompt,
temperature=llm_cfg.temperature,
max_output_tokens=llm_cfg.max_tokens,
max_output_tokens=max_output_tokens,
thinking_config=thinking_config,
response_mime_type="application/json" if json_output else None,
),
)