""" extractor.py - Stage 2: per-sheet construction object extraction. One vision call per page produces a "sheet object manifest": the discipline plus a list of typed, grounded construction objects (each with verbatim source_text and a location_key). Output feeds Stage 3 normalization and Stage 4 clustering. Grounding guard (anti-hallucination): an object is dropped if its primary value contains numeric claims not present in its source_text. Graphical objects (graphical_basis set, no source_text) are allowed through. """ import re import threading from concurrent.futures import ThreadPoolExecutor, as_completed from typing import List, Dict, Optional from backend import config from backend.llm import call_json from backend.prompts import ( EXTRACTOR_SYSTEM_PROMPT, EXTRACTOR_USER_INSTRUCTION, DISCIPLINE_PREFIXES, ATTRIBUTE_VOCAB, ) # prefix (upper) -> discipline, longest-prefix-first for greedy matching _PREFIX_TO_DISCIPLINE = sorted( ((p.upper(), disc) for disc, ps in DISCIPLINE_PREFIXES.items() for p in ps), key=lambda kv: -len(kv[0]), ) _LETTERS_RE = re.compile(r"^[A-Za-z]+") _DIGITS_RE = re.compile(r"\d+") # "Room 124", "RM 124A", "SPACE 12" -> "124" / "124A" / "12". Conservative: an # explicit room/rm/space keyword followed by a 2-4 digit number (+ optional letter). _ROOM_RE = re.compile(r"(?i)\b(?:room|rm|space)\s*#?\s*(\d{2,4}[A-Za-z]?)\b") def _room_from_text(text: Optional[str]) -> Optional[str]: """Best-effort room id from free text when location_key.room is missing.""" if not text: return None m = _ROOM_RE.search(str(text)) return m.group(1) if m else None def discipline_from_sheet_number(sheet_number: Optional[str]) -> Optional[str]: """Deterministic discipline from the sheet-number letter prefix, or None.""" if not sheet_number: return None m = _LETTERS_RE.match(sheet_number.strip()) if not m: return None prefix = m.group(0).upper() for cand, disc in _PREFIX_TO_DISCIPLINE: # longest prefix wins if prefix == cand: return disc # Fall back to the single leading letter (e.g. "A2" -> "A") for cand, disc in _PREFIX_TO_DISCIPLINE: if prefix.startswith(cand) and len(cand) == 1: return disc return None def _is_grounded(value: str, source_text: str, graphical_basis: str = "", page_text: Optional[str] = None) -> bool: """ Keep an object only if its primary value is supported by its source_text, OR it is a graphical object (has graphical_basis with no text to quote). - If graphical_basis is set and source_text is absent, the object is valid. - If the value contains digits, every distinct digit-run must appear in source_text (catches invented dimensions/counts/elevations). - Rescue tier: when page_text (the deterministic text layer) is given, digit-runs absent from source_text but present in the page text are still grounded - vision quoted imperfectly but the value is real. - If the value has no digits, require some alphabetic-token overlap. """ # Graphical objects (no readable text on sheet) are always allowed through. if graphical_basis and not source_text: return True if not value or not source_text: return False value = str(value) src_low = source_text.lower() val_digits = set(_DIGITS_RE.findall(value)) if val_digits: src_digits = set(_DIGITS_RE.findall(source_text)) if val_digits.issubset(src_digits): return True if page_text: return val_digits.issubset(set(_DIGITS_RE.findall(page_text))) return False # No digits: text-based grounding. val_norm = re.sub(r"[^a-z0-9]+", " ", value.lower()).strip() if val_norm and val_norm in src_low: return True val_tokens = {t for t in val_norm.split() if len(t) > 2} src_tokens = set(re.sub(r"[^a-z0-9]+", " ", src_low).split()) return bool(val_tokens & src_tokens) def _primary_value(obj: Dict) -> str: """Extract a single string 'value' from an object for grounding checks. Handles both the new 'attributes' dict schema and the legacy 'value' field. """ attrs = obj.get("attributes") if isinstance(attrs, dict) and attrs: return str(next(iter(attrs.values()))) # Legacy assertion schema has a top-level 'value' field return (obj.get("value") or obj.get("description") or obj.get("name") or obj.get("tag") or "") def _grounding_stamp(value: str, source_text: str, page_text: Optional[str]) -> Optional[str]: """\"text_layer\" when the object survived only via the text-layer rescue tier (digits absent from source_text but present in the page text).""" if not page_text: return None val_digits = set(_DIGITS_RE.findall(str(value))) if not val_digits: return None if val_digits.issubset(set(_DIGITS_RE.findall(source_text))): return None if val_digits.issubset(set(_DIGITS_RE.findall(page_text))): return "text_layer" return None def _normalize_sheet(parsed: Dict, page_number: int, page_text: Optional[str] = None) -> Dict: """ Validate + clean one parsed sheet result, attaching page_number and ids. The LLM now returns a 'sheet' header + 'objects' array (Stage 2 schema). We map each object to an enriched assertion dict that is backward-compatible with all downstream stages (clusterer, conflict_checker, etc.) while also carrying the new typed-object fields. """ # Support both new schema {sheet:{...}, objects:[...]} and old {sheet_number, assertions:[...]} sheet_meta = parsed.get("sheet") or {} sheet_number = sheet_meta.get("sheet_number") or parsed.get("sheet_number") sheet_title = sheet_meta.get("sheet_title") or parsed.get("sheet_title") level = sheet_meta.get("level") or parsed.get("level") scale = sheet_meta.get("scale") or parsed.get("scale") drawing_type = sheet_meta.get("drawing_type") raw_discipline = sheet_meta.get("discipline") or parsed.get("discipline") discipline = raw_discipline or discipline_from_sheet_number(sheet_number) if not discipline or discipline == "General": deduced = discipline_from_sheet_number(sheet_number) if deduced: discipline = deduced discipline = discipline or "Unknown" # Accept both new 'objects' and legacy 'assertions' keys raw_objects = parsed.get("objects") or parsed.get("assertions") or [] clean: List[Dict] = [] dropped = 0 rescued = 0 for idx, obj in enumerate(raw_objects): if not isinstance(obj, dict): continue source_text = obj.get("source_text") or "" graphical_basis = obj.get("graphical_basis") or "" # Derive a primary value for the grounding check primary_val = _primary_value(obj) if not _is_grounded(primary_val, source_text, graphical_basis, page_text=page_text): dropped += 1 continue grounding = _grounding_stamp(primary_val, source_text, page_text) if grounding: rescued += 1 # --- location_key: new schema is richer; map to legacy shape + extras --- lk = obj.get("location_key") if not isinstance(lk, dict): lk = {} # Room recovery: try new keys first, then fall back to text scan room = (lk.get("room") or lk.get("room_number") or _room_from_text(obj.get("name")) or _room_from_text(source_text)) room_name = lk.get("room_name") assertion_id = (obj.get("object_id") or obj.get("assertion_id") or f"{sheet_number or 'p' + str(page_number)}#{idx}") # Map object's attributes dict to legacy subject/attribute/value attrs = obj.get("attributes") or {} if isinstance(attrs, dict) and attrs: first_attr_key = next(iter(attrs)) first_attr_val = str(attrs[first_attr_key]) else: first_attr_key = obj.get("object_type") or "unspecified" first_attr_val = primary_val clean.append({ # --- legacy fields (backward-compat with clusterer, conflict_checker, etc.) --- "id": assertion_id, "subject": (obj.get("name") or obj.get("tag") or obj.get("object_type") or obj.get("subject") or ""), "attribute": first_attr_key, "value": first_attr_val, "location_key": { "grid": lk.get("grid"), "room": room, "room_name": room_name, "level": lk.get("level") or level, "tag": lk.get("tag") or obj.get("tag"), "detail_reference": lk.get("detail_reference"), "plan_zone": lk.get("plan_zone"), "elevation_reference": lk.get("elevation_reference"), }, "source_text": source_text, "confidence": obj.get("confidence") or "medium", # --- new typed-object fields --- "object_type": obj.get("object_type"), "category": obj.get("category"), "object_tag": obj.get("tag"), "object_name": obj.get("name"), "object_description": obj.get("description"), "object_attributes": attrs, "graphical_basis": graphical_basis or None, "review_uses": obj.get("review_uses") or [], **({"grounding": grounding} if grounding else {}), }) if dropped or rescued: print(f"[Extract] Page {page_number} ({sheet_number}): " f"dropped {dropped} ungrounded object(s)" + (f", rescued {rescued} via text layer" if rescued else "")) unresolved = parsed.get("unresolved_items") or [] return { "page_number": page_number, "sheet_number": sheet_number, "discipline": discipline, "sheet_title": sheet_title, "drawing_type": drawing_type, "level": level, "scale": scale, "assertions": clean, # downstream stages read this key "unresolved_items": unresolved, } def _text_layer_block(page: Dict) -> str: """ The TEXT LAYER block appended to the extractor instruction at call sites (NOT a template placeholder - render() silently leaves missing keys as literals). Empty string when the page has no usable text layer. """ text = (page.get("text_layer") or "").strip() if not text: return "" return ("\n\nTEXT LAYER (authoritative for alphanumeric content — trust it " "over the image for numbers, tags, and note text):\n" + text[:config.TEXT_LAYER_MAX_CHARS]) def _extract_one(page: Dict, sheet_hint: str = "") -> Dict: user_text = (EXTRACTOR_USER_INSTRUCTION.replace("{sheet_hint}", sheet_hint) + _text_layer_block(page)) parsed = call_json( system_prompt=EXTRACTOR_SYSTEM_PROMPT, user_text=user_text, images_b64=[page["base64"]], max_tokens=config.EXTRACT_MAX_TOKENS, ) if not isinstance(parsed, dict): return { "page_number": page["page_number"], "sheet_number": None, "discipline": "Unknown", "sheet_title": f"Page {page['page_number']} (extraction failed)", "level": None, "scale": None, "assertions": [], } return _normalize_sheet(parsed, page["page_number"], page_text=page.get("text_layer")) def extract_assertions(pages: List[Dict], on_progress=None) -> List[Dict]: """ Run Stage 2 over all pages. Returns a list of per-sheet object dicts, ordered by page_number. The 'assertions' key on each sheet holds the enriched construction objects (backward-compatible with downstream stages). """ print(f"[Extract] Extracting construction objects from {len(pages)} pages...") results: List[Dict] = [] completed = 0 lock = threading.Lock() def work(page: Dict) -> Dict: nonlocal completed out = _extract_one(page) with lock: completed += 1 n = len(out["assertions"]) u = len(out.get("unresolved_items") or []) print(f"[Extract] Page {page['page_number']} -> {out['discipline']} " f"({out.get('sheet_number')}), {n} objects" + (f", {u} unresolved" if u else "") + f" [{completed}/{len(pages)}]") if on_progress: on_progress(completed, len(pages)) return out with ThreadPoolExecutor(max_workers=config.EXTRACT_CONCURRENCY) as ex: futures = {ex.submit(work, p): p for p in pages} for fut in as_completed(futures): try: results.append(fut.result()) except Exception as e: p = futures[fut] print(f"[Extract] Page {p['page_number']} unexpected error: {e}") results.sort(key=lambda s: s["page_number"]) total = sum(len(s["assertions"]) for s in results) total_unresolved = sum(len(s.get("unresolved_items") or []) for s in results) print(f"[Extract] Done - {total} grounded objects across {len(results)} sheets" + (f" ({total_unresolved} unresolved items)" if total_unresolved else "")) return results