diff --git a/config.yaml b/config.yaml index 55eaec1..288b8a2 100644 --- a/config.yaml +++ b/config.yaml @@ -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 ──────────────────────────────────────────────────────── diff --git a/pipeline/config.py b/pipeline/config.py index 278ec48..217ddc7 100644 --- a/pipeline/config.py +++ b/pipeline/config.py @@ -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( diff --git a/pipeline/stages/llm.py b/pipeline/stages/llm.py index 04ac9f7..8adbc1f 100644 --- a/pipeline/stages/llm.py +++ b/pipeline/stages/llm.py @@ -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, ), )