Kill extract-wave truncation: 65k ceiling, hard thinking budget, reasoning-token telemetry
Job 98194fa8d215 showed every extract call hitting the 32k cap with only ~20k chars visible despite reasoning effort=low - Gemini 2.5 Pro still burned ~25k thinking tokens per sheet. - EXTRACT_MAX_TOKENS default 32768 -> 65536 (model output ceiling) - new EXTRACT_REASONING_MAX_TOKENS (default 2048): OpenRouter reasoning max_tokens / Gemini thinking_budget; takes precedence over effort - log per-call reasoning token counts (usage.completion_tokens_details) and include thinking count in the finish_reason=length marker
This commit is contained in:
@@ -57,6 +57,7 @@ class SheetExtractorAgent:
|
||||
usage_tracker=self.usage,
|
||||
usage_stage="agent.extract",
|
||||
reasoning_effort=config.EXTRACT_REASONING_EFFORT or None,
|
||||
reasoning_max_tokens=config.EXTRACT_REASONING_MAX_TOKENS or None,
|
||||
)
|
||||
|
||||
def run(self, scope: AgentScope) -> AgentResult:
|
||||
|
||||
+10
-3
@@ -74,13 +74,20 @@ MAX_PAGES = int(os.getenv("MAX_PAGES", "60"))
|
||||
MAX_DIMENSION = int(os.getenv("MAX_DIMENSION", "2400")) # px cap on the long edge
|
||||
LLM_TIMEOUT = int(os.getenv("LLM_TIMEOUT", "180")) # seconds per call
|
||||
# Gemini 2.5 Pro counts thinking tokens against max_tokens, so the visible
|
||||
# JSON budget is well under this number on dense sheets. 32768 leaves real
|
||||
# headroom; raise via env if a set still truncates.
|
||||
EXTRACT_MAX_TOKENS = int(os.getenv("EXTRACT_MAX_TOKENS", "32768"))
|
||||
# JSON budget is well under this number on dense sheets. 65536 is the model's
|
||||
# output ceiling - give thinking all the room it wants so visible JSON never
|
||||
# truncates; the thinking budget itself is capped separately below.
|
||||
EXTRACT_MAX_TOKENS = int(os.getenv("EXTRACT_MAX_TOKENS", "65536"))
|
||||
# Reasoning effort for the per-sheet extractor (OpenRouter reasoning knob).
|
||||
# Extraction is perceptive, not deliberative - "low" keeps thinking tokens
|
||||
# from eating the output budget. Empty string disables the parameter.
|
||||
EXTRACT_REASONING_EFFORT = os.getenv("EXTRACT_REASONING_EFFORT", "low").strip()
|
||||
# Hard thinking-token budget for the extractor (OpenRouter reasoning
|
||||
# max_tokens -> Gemini thinking_budget). "low" effort alone still let Gemini
|
||||
# burn ~25k thinking tokens per sheet (job 98194fa8d215); a hard cap forces
|
||||
# the budget into visible output. 0 disables -> falls back to the effort knob.
|
||||
# Mutually exclusive with effort when set (OpenRouter rejects both together).
|
||||
EXTRACT_REASONING_MAX_TOKENS = int(os.getenv("EXTRACT_REASONING_MAX_TOKENS", "2048"))
|
||||
REASON_MAX_TOKENS = int(os.getenv("REASON_MAX_TOKENS", "4096"))
|
||||
|
||||
# -- QAQC stage knobs (Stages 0-1, 3, 6-11) -------------------------
|
||||
|
||||
+39
-11
@@ -117,7 +117,8 @@ def _dump_raw(seq: int, be: Dict[str, Any], usage_stage: str,
|
||||
def _log_call(seq: int, be: Dict[str, Any], usage_stage: str,
|
||||
user_text: str, images_b64: Optional[List[str]],
|
||||
raw: str, parsed: Any, usd: Optional[float],
|
||||
cached: bool = False) -> None:
|
||||
cached: bool = False,
|
||||
reasoning_tokens: Optional[int] = None) -> None:
|
||||
"""One verbose per-call line for the job log (tee'd by job_log.py)."""
|
||||
if not config.LLM_VERBOSE:
|
||||
return
|
||||
@@ -128,9 +129,10 @@ def _log_call(seq: int, be: Dict[str, Any], usage_stage: str,
|
||||
cost_str = f"${usd:.4f}"
|
||||
else:
|
||||
cost_str = "cost n/a"
|
||||
think_str = f"think {reasoning_tokens}tk | " if reasoning_tokens is not None else ""
|
||||
print(f"[LLM] #{seq:04d} {usage_stage} | {be['model']} ({backend}) | "
|
||||
f"in {len(user_text)}ch+{len(images_b64 or [])}img | "
|
||||
f"out {len(raw)}ch | {cost_str} | {_summarize_parsed(parsed)}")
|
||||
f"out {len(raw)}ch | {think_str}{cost_str} | {_summarize_parsed(parsed)}")
|
||||
|
||||
|
||||
def set_text_backend(local: bool) -> None:
|
||||
@@ -230,10 +232,11 @@ def _add_cached() -> None:
|
||||
|
||||
def _cache_key(model: str, system_prompt: str, user_text: str,
|
||||
images_b64: Optional[List[str]], max_tokens: int,
|
||||
json_mode: bool, reasoning_effort: Optional[str] = None) -> str:
|
||||
json_mode: bool, reasoning_effort: Optional[str] = None,
|
||||
reasoning_max_tokens: Optional[int] = None) -> str:
|
||||
h = hashlib.sha256()
|
||||
parts = [model, str(max_tokens), str(json_mode), str(reasoning_effort),
|
||||
system_prompt, user_text]
|
||||
str(reasoning_max_tokens), system_prompt, user_text]
|
||||
for b in (images_b64 or []):
|
||||
parts.append(b)
|
||||
for p in parts:
|
||||
@@ -372,6 +375,21 @@ def _response_cost(response) -> Optional[float]:
|
||||
return float(cost) if isinstance(cost, (int, float)) else None
|
||||
|
||||
|
||||
def _reasoning_tokens(response) -> Optional[int]:
|
||||
"""Hidden thinking tokens for this call (OpenRouter usage details).
|
||||
|
||||
Direct evidence of whether the reasoning knob is working - without it,
|
||||
thinking burn can only be inferred from char counts vs the token cap."""
|
||||
try:
|
||||
dump = response.model_dump()
|
||||
except Exception:
|
||||
return None
|
||||
usage = dump.get("usage") or {}
|
||||
details = usage.get("completion_tokens_details") or {}
|
||||
n = details.get("reasoning_tokens")
|
||||
return int(n) if isinstance(n, (int, float)) else None
|
||||
|
||||
|
||||
def _parse(raw: str) -> Optional[Dict[str, Any]]:
|
||||
"""Parse model JSON, falling back to truncation repair."""
|
||||
try:
|
||||
@@ -392,6 +410,7 @@ def call_json(
|
||||
usage_tracker: Optional[Any] = None,
|
||||
usage_stage: str = "?",
|
||||
reasoning_effort: Optional[str] = None,
|
||||
reasoning_max_tokens: Optional[int] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Send one chat completion expecting a JSON object back.
|
||||
@@ -401,8 +420,10 @@ def call_json(
|
||||
truncation. images_b64: optional base64 JPEGs attached as high-detail image
|
||||
parts. reasoning_effort: optional OpenRouter reasoning knob ("low"/"medium"/
|
||||
"high") - keeps thinking models from burning the output budget on hidden
|
||||
reasoning. Returns the parsed dict, or None on a hard failure (caller
|
||||
degrades).
|
||||
reasoning. reasoning_max_tokens: optional hard thinking-token budget
|
||||
(OpenRouter reasoning max_tokens -> Gemini thinking_budget); stronger than
|
||||
effort, and takes precedence when both are given. Returns the parsed dict,
|
||||
or None on a hard failure (caller degrades).
|
||||
"""
|
||||
has_images = bool(images_b64)
|
||||
be = _resolve_backend(has_images, model)
|
||||
@@ -412,7 +433,8 @@ def call_json(
|
||||
if config.LLM_CACHE:
|
||||
cache_key = _cache_key(be["model"], system_prompt, user_text,
|
||||
images_b64, max_tokens, json_mode=True,
|
||||
reasoning_effort=reasoning_effort)
|
||||
reasoning_effort=reasoning_effort,
|
||||
reasoning_max_tokens=reasoning_max_tokens)
|
||||
hit = _cache_get(cache_key)
|
||||
if hit is not None:
|
||||
_add_cached()
|
||||
@@ -448,8 +470,12 @@ def call_json(
|
||||
if be["usage"]:
|
||||
extra_body["usage"] = {"include": True}
|
||||
# OpenRouter reasoning knob; only sent to cloud backends (local
|
||||
# servers reject unknown fields). Skipped when effort is blank.
|
||||
if reasoning_effort and not be.get("local"):
|
||||
# servers reject unknown fields). A hard thinking budget wins over
|
||||
# the vaguer effort tier - OpenRouter treats them as exclusive.
|
||||
if not be.get("local"):
|
||||
if reasoning_max_tokens:
|
||||
extra_body["reasoning"] = {"max_tokens": int(reasoning_max_tokens)}
|
||||
elif reasoning_effort:
|
||||
extra_body["reasoning"] = {"effort": reasoning_effort}
|
||||
if extra_body:
|
||||
kwargs["extra_body"] = extra_body
|
||||
@@ -466,19 +492,21 @@ def call_json(
|
||||
choice = response.choices[0]
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
raw = _strip_fences(choice.message.content or "")
|
||||
reasoning_tokens = _reasoning_tokens(response)
|
||||
if finish_reason == "length":
|
||||
# Hit max_tokens (thinking tokens included on reasoning models).
|
||||
# Logged explicitly so silent truncation isn't mistaken for a
|
||||
# parse problem; _parse below still salvages what it can.
|
||||
print(f"[LLM] output hit max_tokens (finish_reason=length, "
|
||||
f"{len(raw)}ch returned)")
|
||||
f"{len(raw)}ch returned, "
|
||||
f"thinking={reasoning_tokens}tk)")
|
||||
parsed = _parse(raw)
|
||||
if parsed is not None:
|
||||
_record_model(be, has_images, fell_back)
|
||||
if cache_key:
|
||||
_cache_set(cache_key, parsed)
|
||||
_log_call(seq, be, usage_stage, user_text, images_b64,
|
||||
raw, parsed, usd)
|
||||
raw, parsed, usd, reasoning_tokens=reasoning_tokens)
|
||||
if config.LLM_RAW_DUMP:
|
||||
_dump_raw(seq, be, usage_stage, system_prompt, user_text,
|
||||
images_b64, max_tokens, raw, parsed,
|
||||
|
||||
@@ -50,8 +50,9 @@ def test_run_accepts_bare_list_response():
|
||||
assert len(sheet["assertions"]) == 5
|
||||
# No compact retry needed when the first call yields data.
|
||||
assert mock_call.call_count == 1
|
||||
# Reasoning effort knob is forwarded (None when config is blank in tests).
|
||||
# Reasoning knobs are forwarded (None when config is blank in tests).
|
||||
assert "reasoning_effort" in mock_call.call_args.kwargs
|
||||
assert "reasoning_max_tokens" in mock_call.call_args.kwargs
|
||||
|
||||
|
||||
def test_run_compact_retry_after_hard_failure():
|
||||
|
||||
Reference in New Issue
Block a user