Conflict_Checker/backend/agents/linker.py

"""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 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)
    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)

    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},
            ))
    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}