242 lines
11 KiB
Python
242 lines
11 KiB
Python
"""Scoped extraction and orientation agents."""
|
|
|
|
import json
|
|
from typing import Dict
|
|
|
|
from backend import config
|
|
from backend.agents.base import AgentResult, AgentScope, AgentUsage, failure
|
|
from backend.llm import call_json
|
|
from backend.pipeline.extractor import (
|
|
_normalize_sheet,
|
|
_text_layer_block,
|
|
discipline_from_sheet_number,
|
|
)
|
|
from backend.pipeline.sheet_index import _index_input
|
|
from backend.prompts import (
|
|
EXTRACTOR_SYSTEM_PROMPT,
|
|
EXTRACTOR_USER_INSTRUCTION,
|
|
JURISDICTION_SYSTEM_PROMPT,
|
|
JURISDICTION_USER_INSTRUCTION,
|
|
SHEET_INDEX_SYSTEM_PROMPT,
|
|
SHEET_INDEX_USER_INSTRUCTION,
|
|
)
|
|
|
|
|
|
# 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,
|
|
reasoning_max_tokens=config.EXTRACT_REASONING_MAX_TOKENS or None,
|
|
)
|
|
|
|
def _text_structuring_call(self, instruction_page: Dict, sheet_hint: str):
|
|
"""Rung 2: text-only structuring pass over the page's text layer
|
|
(no image). Recovers text content the vision pass missed."""
|
|
from backend.prompts import (TEXT_STRUCTURING_SYSTEM_PROMPT,
|
|
TEXT_STRUCTURING_USER_INSTRUCTION)
|
|
instruction = (TEXT_STRUCTURING_USER_INSTRUCTION
|
|
.replace("{sheet_hint}", str(sheet_hint or ""))
|
|
.replace("{text_layer}",
|
|
(instruction_page.get("text_layer") or "")
|
|
[:config.TEXT_LAYER_MAX_CHARS]))
|
|
return call_json(
|
|
system_prompt=TEXT_STRUCTURING_SYSTEM_PROMPT,
|
|
user_text=instruction,
|
|
images_b64=None,
|
|
max_tokens=config.EXTRACT_MAX_TOKENS,
|
|
model=config.AGENT_EXTRACT_MODEL,
|
|
usage_tracker=self.usage,
|
|
usage_stage="agent.extract_text",
|
|
reasoning_effort=config.EXTRACT_REASONING_EFFORT or None,
|
|
reasoning_max_tokens=config.EXTRACT_REASONING_MAX_TOKENS or None,
|
|
)
|
|
|
|
def run(self, scope: AgentScope) -> AgentResult:
|
|
from backend.text_coverage import (fallback_objects, merge_objects,
|
|
recover_sheet_number, text_coverage)
|
|
try:
|
|
page = scope.payload["page"]
|
|
hint = scope.payload.get("sheet_hint") or ""
|
|
page_text = page.get("text_layer")
|
|
instruction = EXTRACTOR_USER_INSTRUCTION.replace(
|
|
"{sheet_hint}", str(hint)) + _text_layer_block(page)
|
|
|
|
# Rung 1: vision pass (unchanged behaviour, incl. compact retry)
|
|
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):
|
|
# Don't give up on the page - the ladder below can still
|
|
# rescue it from the text layer.
|
|
parsed = {"sheet": {}, "objects": []}
|
|
|
|
sheet = _normalize_sheet(parsed, page["page_number"],
|
|
page_text=page_text)
|
|
cov = text_coverage(page_text or "", sheet["assertions"])
|
|
sheet["coverage"] = cov
|
|
|
|
# Rung 2: text-only structuring when coverage is below floor.
|
|
# MERGE, never replace: vision keeps every object it found
|
|
# (graphical_basis content exists only in the image); the text
|
|
# pass fills in the text content the vision pass missed.
|
|
if (page_text and config.EXTRACT_TEXT_RETRY_ENABLED
|
|
and cov["ratio"] < config.EXTRACT_COVERAGE_FLOOR):
|
|
print(f"[Extract] Page {page['page_number']}: coverage "
|
|
f"{cov['ratio']:.0%} < floor - text-only structuring pass")
|
|
parsed2 = _wrap_bare_list(
|
|
self._text_structuring_call(page, hint), page["page_number"])
|
|
if isinstance(parsed2, dict):
|
|
sheet2 = _normalize_sheet(parsed2, page["page_number"],
|
|
page_text=page_text)
|
|
before = len(sheet["assertions"])
|
|
sheet["assertions"] = merge_objects(sheet["assertions"],
|
|
sheet2["assertions"])
|
|
# Fill header gaps the vision pass left null
|
|
for key in ("sheet_number", "sheet_title", "discipline",
|
|
"level", "scale", "drawing_type"):
|
|
if not sheet.get(key) and sheet2.get(key):
|
|
sheet[key] = sheet2[key]
|
|
cov = text_coverage(page_text, sheet["assertions"])
|
|
sheet["coverage"] = cov
|
|
print(f"[Extract] Page {page['page_number']}: merged "
|
|
f"{len(sheet['assertions']) - before} text-structured "
|
|
f"object(s), coverage now {cov['ratio']:.0%}")
|
|
|
|
# Rung 3: deterministic fallback - dark sheets are impossible.
|
|
# Also merged (deduped) so stub notes never double up with
|
|
# objects the earlier rungs already captured.
|
|
if (page_text and config.EXTRACT_FALLBACK_ENABLED
|
|
and cov["ratio"] < config.EXTRACT_COVERAGE_FLOOR):
|
|
stubs = fallback_objects(page_text, page["page_number"],
|
|
config.EXTRACT_FALLBACK_MAX_OBJECTS)
|
|
stubs = _normalize_sheet({"sheet": {}, "objects": stubs},
|
|
page["page_number"],
|
|
page_text=page_text)["assertions"]
|
|
# _normalize_sheet only stamps its own "text_layer" rescue
|
|
# grounding; restore the explicit fallback provenance.
|
|
for stub in stubs:
|
|
stub["grounding"] = "text_layer_fallback"
|
|
before = len(sheet["assertions"])
|
|
sheet["assertions"] = merge_objects(sheet["assertions"], stubs)
|
|
print(f"[Extract] Page {page['page_number']}: fallback merged "
|
|
f"{len(sheet['assertions']) - before} text-layer stub(s)")
|
|
sheet["coverage"] = text_coverage(page_text,
|
|
sheet["assertions"])
|
|
|
|
# Identity recovery: never leave a text-bearing page sheet-less
|
|
if not sheet.get("sheet_number") and page_text:
|
|
recovered = recover_sheet_number(page_text)
|
|
if recovered:
|
|
sheet["sheet_number"] = recovered
|
|
sheet["discipline"] = (
|
|
discipline_from_sheet_number(recovered)
|
|
or sheet.get("discipline") or "Unknown")
|
|
print(f"[Extract] Page {page['page_number']}: sheet number "
|
|
f"recovered from text layer -> {recovered}")
|
|
|
|
return AgentResult(scope_id=scope.scope_id, artifacts=[sheet])
|
|
except Exception as exc:
|
|
return failure(scope, exc)
|
|
|
|
|
|
class SheetIndexAgent:
|
|
name = "sheet_index"
|
|
|
|
def __init__(self, usage: AgentUsage) -> None:
|
|
self.usage = usage
|
|
|
|
def run(self, scope: AgentScope) -> AgentResult:
|
|
try:
|
|
sheets = scope.payload.get("sheets") or []
|
|
instruction = SHEET_INDEX_USER_INSTRUCTION.replace(
|
|
"{sheet_index_input}",
|
|
json.dumps(_index_input(sheets), ensure_ascii=True),
|
|
)
|
|
parsed = call_json(
|
|
system_prompt=SHEET_INDEX_SYSTEM_PROMPT,
|
|
user_text=instruction,
|
|
max_tokens=config.SHEET_INDEX_MAX_TOKENS,
|
|
model=config.AGENT_INDEX_MODEL,
|
|
usage_tracker=self.usage,
|
|
usage_stage="agent.sheet_index",
|
|
)
|
|
if isinstance(parsed, list):
|
|
parsed = {"sheet_index": parsed, "missing_expected_sheets": []}
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError("no sheet index returned")
|
|
return AgentResult(scope_id=scope.scope_id, artifacts=[parsed])
|
|
except Exception as exc:
|
|
return failure(scope, exc)
|
|
|
|
|
|
class JurisdictionAgent:
|
|
name = "jurisdiction"
|
|
|
|
def __init__(self, usage: AgentUsage) -> None:
|
|
self.usage = usage
|
|
|
|
def run(self, scope: AgentScope) -> AgentResult:
|
|
try:
|
|
project_input: Dict = scope.payload.get("project_input") or {}
|
|
instruction = JURISDICTION_USER_INSTRUCTION.replace(
|
|
"{project_input}", json.dumps(project_input, ensure_ascii=True)
|
|
)
|
|
parsed = call_json(
|
|
system_prompt=JURISDICTION_SYSTEM_PROMPT,
|
|
user_text=instruction,
|
|
max_tokens=config.JURISDICTION_MAX_TOKENS,
|
|
model=config.AGENT_JURISDICTION_MODEL,
|
|
usage_tracker=self.usage,
|
|
usage_stage="agent.jurisdiction",
|
|
)
|
|
if not isinstance(parsed, dict):
|
|
raise ValueError("no jurisdiction profile returned")
|
|
profile = parsed.get("project_code_profile")
|
|
artifact = profile if isinstance(profile, dict) else parsed
|
|
return AgentResult(scope_id=scope.scope_id, artifacts=[artifact])
|
|
except Exception as exc:
|
|
return failure(scope, exc)
|