""" normalizer.py - Stage 3: construction object normalization + project intelligence. Stage 3a: normalize_assertions — adds normalized values (e.g. 9'-0" -> 108 in) and normalized tags/room identifiers to extracted objects. Additive; originals are always preserved. Processed in batches across a thread pool. Stage 3b: build_project_intelligence — cross-sheet merging pass that assigns Global Object IDs (GOIDs) and builds relationships between objects that represent the same physical element across multiple drawings/schedules. Batched by object_type so each call stays within the token budget. """ import json from collections import defaultdict from concurrent.futures import ThreadPoolExecutor from typing import Dict, List from backend import config from backend.pipeline._stage import call_stage from backend.prompts import ( NORMALIZATION_SYSTEM_PROMPT, NORMALIZATION_USER_INSTRUCTION, PROJECT_INTELLIGENCE_SYSTEM_PROMPT, PROJECT_INTELLIGENCE_USER_INSTRUCTION, ) def _batches(items: List[Dict], size: int) -> List[List[Dict]]: return [items[i:i + size] for i in range(0, len(items), size)] def _normalize_batch(batch: List[Dict]) -> Dict[str, Dict]: """Return {assertion_id: normalization fields} for one batch.""" payload = [ { "assertion_id": a["id"], "attribute": a.get("attribute"), "value": a.get("value"), "object_type": a.get("object_type"), "object_tag": a.get("object_tag"), "source_text": a.get("source_text"), } for a in batch ] parsed = call_stage( NORMALIZATION_SYSTEM_PROMPT, NORMALIZATION_USER_INSTRUCTION, subs={"assertions": json.dumps(payload, ensure_ascii=True)}, max_tokens=config.NORMALIZE_MAX_TOKENS, ) if isinstance(parsed, list): rows = parsed elif isinstance(parsed, dict): rows = parsed.get("normalized_assertions") or [] else: rows = [] out: Dict[str, Dict] = {} for n in rows: if isinstance(n, dict) and n.get("assertion_id"): out[n["assertion_id"]] = { "normalized_value": n.get("normalized_value"), "normalized_unit": n.get("normalized_unit"), "normalized_tag": n.get("normalized_tag"), "normalized_room": n.get("normalized_room"), } return out def normalize_assertions(sheets: List[Dict]) -> List[Dict]: """Annotate each assertion in-place with normalized fields.""" all_assertions = [a for s in sheets for a in s.get("assertions", [])] if not all_assertions: return sheets batches = _batches(all_assertions, config.NORMALIZE_BATCH_SIZE) merged: Dict[str, Dict] = {} with ThreadPoolExecutor(max_workers=config.NORMALIZE_CONCURRENCY) as pool: for result in pool.map(_normalize_batch, batches): merged.update(result) applied = 0 for a in all_assertions: norm = merged.get(a["id"]) if norm: if norm.get("normalized_value") is not None: a["normalized_value"] = norm["normalized_value"] a["normalized_unit"] = norm.get("normalized_unit") applied += 1 if norm.get("normalized_tag"): a["normalized_tag"] = norm["normalized_tag"] if norm.get("normalized_room"): a["normalized_room"] = norm["normalized_room"] # Also push into location_key so clusterer picks it up if isinstance(a.get("location_key"), dict) and not a["location_key"].get("room"): a["location_key"]["room"] = norm["normalized_room"] print(f"[Normalize] normalized {applied}/{len(all_assertions)} objects") return sheets # --------------------------------------------------------------------------- # Stage 3b - project intelligence: GOIDs + relationships # --------------------------------------------------------------------------- # Max objects of one type to send in a single intelligence call. _INTEL_BATCH_SIZE = 60 def _intel_batch(object_type: str, objects: List[Dict]) -> Dict: """Run the project intelligence pass for one object_type batch.""" payload = [] for obj in objects: payload.append({ "object_id": obj.get("id"), "object_type": obj.get("object_type") or object_type, "tag": obj.get("object_tag"), "name": obj.get("object_name"), "attributes": obj.get("object_attributes") or {}, "location_key": obj.get("location_key") or {}, "source_text": obj.get("source_text") or "", "source_sheet": (obj.get("location_key") or {}).get("sheet_number"), }) parsed = call_stage( PROJECT_INTELLIGENCE_SYSTEM_PROMPT, PROJECT_INTELLIGENCE_USER_INSTRUCTION, subs={ "object_type": object_type, "objects": json.dumps(payload, ensure_ascii=True), }, max_tokens=config.NORMALIZE_MAX_TOKENS, ) if not isinstance(parsed, dict): return {} return { "project_objects": parsed.get("project_objects") or [], "unresolved_relationships": parsed.get("unresolved_relationships") or [], } def build_project_intelligence(sheets: List[Dict]) -> Dict: """ Stage 3b: cross-sheet GOID assignment and relationship building. Groups all extracted objects by object_type, then calls the Project Intelligence LLM once per type (batched) to assign GOIDs and link objects that represent the same physical element across sheets. Returns {"project_objects": [...], "unresolved_relationships": [...]}. On any failure, returns an empty dict (graceful degradation). """ all_assertions = [a for s in sheets for a in s.get("assertions", [])] if not all_assertions: return {} # Group by object_type; unknown/graphical go under "general" by_type: Dict[str, List[Dict]] = defaultdict(list) for a in all_assertions: otype = a.get("object_type") or "general" by_type[otype].append(a) all_project_objects: List[Dict] = [] all_unresolved: List[Dict] = [] # Process each type in batches; high-count types (e.g. dimensions) are # less useful for GOID merging so we skip types with only 1 object. for otype, objs in sorted(by_type.items()): if len(objs) < 2: continue batches = _batches(objs, _INTEL_BATCH_SIZE) for batch in batches: try: result = _intel_batch(otype, batch) all_project_objects.extend(result.get("project_objects") or []) all_unresolved.extend(result.get("unresolved_relationships") or []) except Exception as e: print(f"[ProjectIntel] {otype} batch failed: {e}") print(f"[ProjectIntel] {len(all_project_objects)} project objects, " f"{len(all_unresolved)} unresolved relationships") return { "project_objects": all_project_objects, "unresolved_relationships": all_unresolved, }