diff --git a/backend/.env.example b/backend/.env.example index ec5335e..abb66e4 100644 --- a/backend/.env.example +++ b/backend/.env.example @@ -39,7 +39,10 @@ PDF_DPI=100 MAX_PAGES=60 MAX_DIMENSION=2400 LLM_TIMEOUT=180 -EXTRACT_MAX_TOKENS=8192 +EXTRACT_MAX_TOKENS=32768 +# Reasoning effort for per-sheet extraction (low keeps Gemini thinking tokens +# from eating the output budget). Blank = don't send the parameter. +EXTRACT_REASONING_EFFORT=low REASON_MAX_TOKENS=4096 EXTRACT_CONCURRENCY=4 REASON_CONCURRENCY=4 diff --git a/backend/agents/extractors.py b/backend/agents/extractors.py index 2123a90..eb41afa 100644 --- a/backend/agents/extractors.py +++ b/backend/agents/extractors.py @@ -18,27 +18,64 @@ from backend.prompts import ( ) +# Appended to the extractor instruction on the second-chance retry. Dense plan +# sheets blow the output budget on the full schema; compact mode trades +# per-object verbosity for actually finishing the page. +_COMPACT_RETRY_SUFFIX = """ +IMPORTANT - COMPACT RETRY: the first pass did not complete. Keep the SAME JSON +schema, but extract at most 40 objects, prioritizing coordination-relevant +items (equipment, fixtures, devices, keynotes, dimensions, markers/callouts, +schedule rows). Keep descriptions/attributes short; skip review_uses entries +you are unsure about. Finish the JSON - a smaller complete answer beats a +larger truncated one.""" + + +def _wrap_bare_list(parsed, page_number: int): + """Models sometimes skip the {sheet, objects} wrapper and return a bare + objects array (especially after truncation repair). Accept it - the sheet + header falls back to title-block deduction downstream.""" + if isinstance(parsed, list): + print(f"[Extract] Page {page_number}: wrapping bare objects array " + f"({len(parsed)} items, no sheet header)") + return {"sheet": {}, "objects": parsed} + return parsed + + class SheetExtractorAgent: name = "sheet_extractor" def __init__(self, usage: AgentUsage) -> None: self.usage = usage + def _call(self, instruction: str, page: Dict): + return call_json( + system_prompt=EXTRACTOR_SYSTEM_PROMPT, + user_text=instruction, + images_b64=[page["base64"]], + max_tokens=config.EXTRACT_MAX_TOKENS, + model=config.AGENT_EXTRACT_MODEL, + usage_tracker=self.usage, + usage_stage="agent.extract", + reasoning_effort=config.EXTRACT_REASONING_EFFORT or None, + ) + def run(self, scope: AgentScope) -> AgentResult: try: page = scope.payload["page"] instruction = EXTRACTOR_USER_INSTRUCTION.replace( "{sheet_hint}", str(scope.payload.get("sheet_hint") or "") ) - parsed = call_json( - system_prompt=EXTRACTOR_SYSTEM_PROMPT, - user_text=instruction, - images_b64=[page["base64"]], - max_tokens=config.EXTRACT_MAX_TOKENS, - model=config.AGENT_EXTRACT_MODEL, - usage_tracker=self.usage, - usage_stage="agent.extract", - ) + parsed = _wrap_bare_list(self._call(instruction, page), + page["page_number"]) + if not isinstance(parsed, dict): + # Second chance: same page, compact instructions. Runs only + # when the full-schema pass returned nothing usable. + print(f"[Extract] Page {page['page_number']}: full extraction " + f"failed, retrying compact") + parsed = _wrap_bare_list( + self._call(instruction + _COMPACT_RETRY_SUFFIX, page), + page["page_number"], + ) if not isinstance(parsed, dict): raise ValueError("no structured extraction returned") sheet = _normalize_sheet(parsed, page["page_number"]) diff --git a/backend/config.py b/backend/config.py index 73c68a9..3d7e9e4 100644 --- a/backend/config.py +++ b/backend/config.py @@ -73,7 +73,14 @@ PDF_DPI = int(os.getenv("PDF_DPI", "100")) 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 -EXTRACT_MAX_TOKENS = int(os.getenv("EXTRACT_MAX_TOKENS", "16384")) +# 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")) +# 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() REASON_MAX_TOKENS = int(os.getenv("REASON_MAX_TOKENS", "4096")) # -- QAQC stage knobs (Stages 0-1, 3, 6-11) ------------------------- diff --git a/backend/llm.py b/backend/llm.py index 3a5cf98..afc722a 100644 --- a/backend/llm.py +++ b/backend/llm.py @@ -81,7 +81,8 @@ def _summarize_parsed(parsed: Any) -> str: def _dump_raw(seq: int, be: Dict[str, Any], usage_stage: str, system_prompt: str, user_text: str, images_b64: Optional[List[str]], max_tokens: int, - raw: str, parsed: Any) -> None: + raw: str, parsed: Any, + finish_reason: Optional[str] = None) -> None: """Write the full request/response for one call to the job's llm_raw dir.""" if not _raw_dump_dir: return @@ -100,6 +101,7 @@ def _dump_raw(seq: int, be: Dict[str, Any], usage_stage: str, "n_images": len(images_b64 or []), "system_prompt": system_prompt, "user_text": user_text, + "finish_reason": finish_reason, "raw_response": raw, "parsed": parsed, } @@ -228,9 +230,10 @@ 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) -> str: + json_mode: bool, reasoning_effort: Optional[str] = None) -> str: h = hashlib.sha256() - parts = [model, str(max_tokens), str(json_mode), system_prompt, user_text] + parts = [model, str(max_tokens), str(json_mode), str(reasoning_effort), + system_prompt, user_text] for b in (images_b64 or []): parts.append(b) for p in parts: @@ -388,6 +391,7 @@ def call_json( model: Optional[str] = None, usage_tracker: Optional[Any] = None, usage_stage: str = "?", + reasoning_effort: Optional[str] = None, ) -> Optional[Dict[str, Any]]: """ Send one chat completion expecting a JSON object back. @@ -395,7 +399,10 @@ def call_json( Uses the provider's JSON mode (response_format) so the model returns a bare JSON object instead of prose/empty text, and recovers from max_tokens truncation. images_b64: optional base64 JPEGs attached as high-detail image - parts. Returns the parsed dict, or None on a hard failure (caller degrades). + 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). """ has_images = bool(images_b64) be = _resolve_backend(has_images, model) @@ -404,7 +411,8 @@ def call_json( cache_key = None if config.LLM_CACHE: cache_key = _cache_key(be["model"], system_prompt, user_text, - images_b64, max_tokens, json_mode=True) + images_b64, max_tokens, json_mode=True, + reasoning_effort=reasoning_effort) hit = _cache_get(cache_key) if hit is not None: _add_cached() @@ -434,10 +442,17 @@ def call_json( for attempt in range(2): try: client = get_client(be["base_url"], be["api_key"]) - kwargs = dict(model=be["model"], messages=messages, - max_tokens=max_tokens, timeout=config.LLM_TIMEOUT) + kwargs: Dict[str, Any] = dict(model=be["model"], messages=messages, + max_tokens=max_tokens, timeout=config.LLM_TIMEOUT) + extra_body: Dict[str, Any] = {} if be["usage"]: - kwargs["extra_body"] = {"usage": {"include": True}} + 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"): + extra_body["reasoning"] = {"effort": reasoning_effort} + if extra_body: + kwargs["extra_body"] = extra_body if use_json_mode: kwargs["response_format"] = {"type": "json_object"} response = client.chat.completions.create(**kwargs) @@ -448,7 +463,15 @@ def call_json( usage_tracker.record( usage_stage, be["model"], usd=usd or 0.0, has_images=has_images ) - raw = _strip_fences(response.choices[0].message.content or "") + choice = response.choices[0] + finish_reason = getattr(choice, "finish_reason", None) + raw = _strip_fences(choice.message.content or "") + 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)") parsed = _parse(raw) if parsed is not None: _record_model(be, has_images, fell_back) @@ -458,7 +481,8 @@ def call_json( raw, parsed, usd) if config.LLM_RAW_DUMP: _dump_raw(seq, be, usage_stage, system_prompt, user_text, - images_b64, max_tokens, raw, parsed) + images_b64, max_tokens, raw, parsed, + finish_reason=finish_reason) return parsed if attempt == 0: print("[LLM] JSON parse error (retrying)") diff --git a/tests/agents/test_sheet_extractor_fallback.py b/tests/agents/test_sheet_extractor_fallback.py new file mode 100644 index 0000000..0c0db37 --- /dev/null +++ b/tests/agents/test_sheet_extractor_fallback.py @@ -0,0 +1,75 @@ +"""SheetExtractorAgent fallback ladder tests (bare-list wrap + compact retry).""" + +from unittest.mock import patch + +from backend.agents.base import AgentScope, AgentUsage +from backend.agents.extractors import SheetExtractorAgent, _wrap_bare_list + + +def _scope(): + return AgentScope( + scope_id="sheet:4", + payload={"page": {"page_number": 4, "base64": "QUJD"}}, + ) + + +def _objects(n=2): + return [ + { + "object_id": f"obj-{i}", + "object_type": "equipment", + "category": "mechanical", + "name": f"RTU-{i}", + "source_text": f"RTU-{i}", + "confidence": "high", + } + for i in range(n) + ] + + +def test_wrap_bare_list_builds_sheet_envelope(): + wrapped = _wrap_bare_list(_objects(3), page_number=4) + assert wrapped["sheet"] == {} + assert len(wrapped["objects"]) == 3 + + +def test_wrap_bare_list_passes_dicts_and_none_through(): + assert _wrap_bare_list({"sheet": {}, "objects": []}, 1) == {"sheet": {}, "objects": []} + assert _wrap_bare_list(None, 1) is None + + +def test_run_accepts_bare_list_response(): + agent = SheetExtractorAgent(usage=AgentUsage()) + with patch("backend.agents.extractors.call_json", + return_value=_objects(5)) as mock_call: + result = agent.run(_scope()) + assert not result.error + assert len(result.artifacts) == 1 + sheet = result.artifacts[0] + assert sheet["page_number"] == 4 + 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). + assert "reasoning_effort" in mock_call.call_args.kwargs + + +def test_run_compact_retry_after_hard_failure(): + agent = SheetExtractorAgent(usage=AgentUsage()) + with patch("backend.agents.extractors.call_json", + side_effect=[None, {"sheet": {"sheet_number": "A102"}, + "objects": _objects(2)}]) as mock_call: + result = agent.run(_scope()) + assert not result.error + assert result.artifacts[0]["sheet_number"] == "A102" + assert mock_call.call_count == 2 + # Second call carried the compact suffix. + assert "COMPACT RETRY" in mock_call.call_args_list[1].kwargs["user_text"] + + +def test_run_fails_only_after_both_attempts_miss(): + agent = SheetExtractorAgent(usage=AgentUsage()) + with patch("backend.agents.extractors.call_json", return_value=None) as mock_call: + result = agent.run(_scope()) + assert result.error == "no structured extraction returned" + assert mock_call.call_count == 2