"""Bounded semantic linkers that build coordination clusters.""" import json import re from collections import defaultdict from typing import Dict, Iterable, List, Tuple from backend import config from backend.agents.base import AgentResult, AgentScope, AgentUsage, failure from backend.llm import call_json from backend.pipeline.clusterer import cluster_by_location from backend.pipeline.llm_clusterer import _location from backend.prompts import CLUSTER_SYSTEM_PROMPT, CLUSTER_USER_INSTRUCTION def _family(assertion: Dict) -> str: location = assertion.get("location_key") or {} if location.get("room"): return "room" if location.get("grid"): return "grid" if location.get("detail_reference"): return "detail" tag = str(location.get("tag") or "") match = re.match(r"[A-Za-z]+", tag) return ( (match.group(0).lower() if match else "") or (assertion.get("object_type") or "").lower() or "general" ) def _xref_keys(assertion: Dict) -> List[str]: """Cross-level join keys: detail references and member tags.""" location = assertion.get("location_key") or {} keys = [] ref = re.sub(r"\s+", "", str(location.get("detail_reference") or "")).upper() if ref: keys.append(f"detail:{ref}") tag = re.sub(r"\s+", "", str(location.get("tag") or "")).upper() if re.match(r"^[A-Z]+\d", tag): # member marks: W12X26, HSS16X4X5/8, ... keys.append(f"tag:{tag}") return keys def build_link_scopes(sheets: List[Dict]) -> List[AgentScope]: """Partition facts by level and object/tag family, then enforce a hard cap.""" buckets: Dict[Tuple[str, str], List[Dict]] = defaultdict(list) xref: Dict[str, List[Dict]] = defaultdict(list) for sheet in sheets: for assertion in sheet.get("assertions", []): enriched = { **assertion, "discipline": sheet.get("discipline") or "Unknown", "sheet_number": sheet.get("sheet_number"), "page_number": sheet.get("page_number"), } level = str((assertion.get("location_key") or {}).get("level") or sheet.get("level") or "unknown").lower() buckets[(level, _family(assertion))].append(enriched) for key in _xref_keys(assertion): xref[key].append(enriched) scopes: List[AgentScope] = [] cap = max(2, config.AGENT_LINK_MAX_ASSERTIONS) for (level, family), assertions in sorted(buckets.items()): for offset in range(0, len(assertions), cap): chunk = assertions[offset:offset + cap] if len(chunk) < 2: continue scopes.append(AgentScope( scope_id=f"{level}:{family}:{offset // cap + 1}", payload={"assertions": chunk, "level": level, "family": family}, )) for key, assertions in sorted(xref.items()): sheets_present = {a.get("sheet_number") for a in assertions} if len(assertions) < 2 or len(sheets_present) < 2: continue chunk = assertions[:cap] if len({a.get("sheet_number") for a in chunk}) < 2: continue # cap landed on a single sheet — xref adds nothing scopes.append(AgentScope( scope_id=f"xref:{key}", payload={"assertions": chunk, "level": "xref", "family": key}, )) return scopes def _payload(assertions: Iterable[Dict]) -> List[Dict]: return [{ "assertion_id": item.get("id"), "discipline": item.get("discipline"), "sheet_number": item.get("sheet_number"), "attribute": item.get("attribute"), "value": item.get("value"), "location_key": item.get("location_key"), "source_text": item.get("source_text"), } for item in assertions] def _fallback(assertions: List[Dict]) -> List[Dict]: """Use the deterministic linker within this scope when semantic linking fails.""" by_sheet: Dict[Tuple, Dict] = {} for item in assertions: key = (item.get("sheet_number"), item.get("page_number")) sheet = by_sheet.setdefault(key, { "sheet_number": item.get("sheet_number"), "page_number": item.get("page_number"), "discipline": item.get("discipline"), "assertions": [], }) sheet["assertions"].append(item) return cluster_by_location(list(by_sheet.values())) class LinkerAgent: name = "linker" def __init__(self, usage: AgentUsage) -> None: self.usage = usage def run(self, scope: AgentScope) -> AgentResult: try: assertions = scope.payload.get("assertions") or [] by_id = {item.get("id"): item for item in assertions if item.get("id")} instruction = CLUSTER_USER_INSTRUCTION.replace( "{normalized_assertions}", json.dumps(_payload(assertions), ensure_ascii=True), ) parsed = call_json( system_prompt=CLUSTER_SYSTEM_PROMPT, user_text=instruction, max_tokens=config.CLUSTER_MAX_TOKENS, model=config.AGENT_LINKER_MODEL, usage_tracker=self.usage, usage_stage="agent.link", ) raw = parsed if isinstance(parsed, list) else ( parsed.get("clusters") if isinstance(parsed, dict) else [] ) clusters: List[Dict] = [] for candidate in raw or []: if not isinstance(candidate, dict): continue members = [ by_id[item_id] for item_id in candidate.get("assertion_ids") or [] if item_id in by_id ] if len(members) < 2: continue allowed_sheets = [] for member in members: sheet = member.get("sheet_number") if sheet not in allowed_sheets: allowed_sheets.append(sheet) allowed_sheets = allowed_sheets[:config.AGENT_CONFLICT_MAX_IMAGES] members = [ member for member in members if member.get("sheet_number") in allowed_sheets ][:config.AGENT_CLUSTER_MAX_ASSERTIONS] primary = candidate.get("primary_location_key") or {} clusters.append({ "key": f"{scope.scope_id}:{candidate.get('cluster_id') or len(clusters) + 1}", "location": _location(primary), "disciplines": sorted({ member.get("discipline") or "Unknown" for member in members }), "page_numbers": sorted({ member["page_number"] for member in members if member.get("page_number") }), "sheets": sorted({ member["sheet_number"] for member in members if member.get("sheet_number") }), "assertions": members, "kind": "agent_semantic", "scope_id": scope.scope_id, }) if not clusters: clusters = _fallback(assertions) for cluster in clusters: cluster["scope_id"] = scope.scope_id cluster["kind"] = "agent_deterministic" return AgentResult(scope_id=scope.scope_id, artifacts=clusters) except Exception as exc: return failure(scope, exc) def build_object_graph(clusters: List[Dict]) -> Dict: """Build a deterministic graph view from linker output.""" nodes = [] edges = [] seen = set() for cluster in clusters: cluster_id = cluster.get("key") nodes.append({ "id": cluster_id, "type": "cluster", "location": cluster.get("location"), "sheets": cluster.get("sheets") or [], }) for assertion in cluster.get("assertions") or []: assertion_id = assertion.get("id") if not assertion_id: continue if assertion_id not in seen: seen.add(assertion_id) nodes.append({ "id": assertion_id, "type": assertion.get("object_type") or "assertion", "sheet": assertion.get("sheet_number"), }) edges.append({ "source": assertion_id, "target": cluster_id, "relationship": "member_of", }) return {"nodes": nodes, "edges": edges}