434 lines
18 KiB
Python
434 lines
18 KiB
Python
"""
|
|
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,
|
|
TEXT_STRUCTURING_SYSTEM_PROMPT,
|
|
TEXT_STRUCTURING_USER_INSTRUCTION,
|
|
DISCIPLINE_PREFIXES,
|
|
ATTRIBUTE_VOCAB,
|
|
)
|
|
from backend.text_coverage import (
|
|
_norm,
|
|
fallback_objects,
|
|
merge_objects,
|
|
recover_sheet_number,
|
|
text_coverage,
|
|
)
|
|
|
|
# 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
|
|
unverified = 0
|
|
page_norm = _norm(page_text) if page_text else ""
|
|
|
|
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
|
|
# Pre-set stamps (fallback/merge rungs) win; otherwise compute the
|
|
# text-layer rescue stamp.
|
|
grounding = obj.get("grounding") or _grounding_stamp(
|
|
primary_val, source_text, page_text)
|
|
if grounding == "text_layer":
|
|
rescued += 1
|
|
if not grounding and page_text and source_text:
|
|
# Vision-unverified: survived the digit guard, but the quoted
|
|
# source_text is not present in the deterministic text layer.
|
|
# Kept and stamped - the wave-5b verifier prioritizes these.
|
|
if _norm(str(source_text)) not in page_norm:
|
|
grounding = "vision_unverified"
|
|
unverified += 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 or unverified:
|
|
print(f"[Extract] Page {page_number} ({sheet_number}): "
|
|
f"dropped {dropped} ungrounded object(s)"
|
|
+ (f", rescued {rescued} via text layer" if rescued else "")
|
|
+ (f", {unverified} vision-unverified" if unverified 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 _text_structuring_extract(page: Dict, sheet_hint: str = ""):
|
|
"""Rung 2 of the extraction ladder: text-only structuring call (no
|
|
image). The text layer is authoritative for alphanumeric content - the
|
|
model segments it instead of transcribing pixels, so vision misreads
|
|
are impossible on this rung."""
|
|
instruction = (TEXT_STRUCTURING_USER_INSTRUCTION
|
|
.replace("{sheet_hint}", sheet_hint or "")
|
|
.replace("{text_layer}",
|
|
(page.get("text_layer") or "")
|
|
[:config.TEXT_LAYER_MAX_CHARS]))
|
|
return call_json(
|
|
system_prompt=TEXT_STRUCTURING_SYSTEM_PROMPT,
|
|
user_text=instruction,
|
|
max_tokens=config.EXTRACT_MAX_TOKENS,
|
|
)
|
|
|
|
|
|
def _extract_one(page: Dict, sheet_hint: str = "") -> Dict:
|
|
page_text = page.get("text_layer")
|
|
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):
|
|
if not page_text:
|
|
# Scanned/raster page: vision-only, keep the legacy failure shape.
|
|
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": [],
|
|
}
|
|
# Text-bearing page: climb the ladder instead of going dark.
|
|
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 objects (graphical_basis content exists only
|
|
# in the image) are kept; the text pass fills what vision 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 = _text_structuring_extract(page, sheet_hint)
|
|
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"])
|
|
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 - a dark text-bearing sheet is
|
|
# impossible. Stubs are deduped against earlier rungs.
|
|
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"]
|
|
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 sheet
|
|
|
|
|
|
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
|