Conflict_Checker/backend/agents/runner.py

"""Public entry point for the scoped Agent-mode pipeline."""

import base64
import json
import os
from typing import Callable, Dict, Optional

from backend import config
from backend.agents.base import AgentScope, AgentUsage
from backend.agents.brain import BrainAgent
from backend.agents.code_agent import CodeAgent, build_code_scopes
from backend.agents.completeness import CompletenessAgent, build_sheet_summaries
from backend.agents.conflict_critic import ConflictCriticAgent
from backend.agents.construct_agent import ConstructabilityAgent, build_construct_scopes
from backend.agents.disputes import annotate_clusters
from backend.agents.extractors import (
    JurisdictionAgent,
    SheetExtractorAgent,
    SheetIndexAgent,
)
from backend.agents.linker import LinkerAgent, build_link_scopes, build_object_graph
from backend.agents.memory import ProjectMemory
from backend.agents.orchestrator import Orchestrator
from backend.agents.rfi_writer import RFIWriterAgent
from backend.agents.verifier import (
    EvidenceVerifierAgent, apply_verdicts, select_findings,
)
from backend.llm import reset_cost
from backend.pipeline.pdf_processor import convert_pdf_to_images
from backend.pipeline.report import build_report, to_markdown
from backend.pipeline.sheet_index import derive_project_meta_from_cover
from backend.review.gate import build_review_queue
from backend.review.store import ReviewStore
from backend.text_layer import (
    attach_text_layers, coverage_gaps, find_evidence_bbox, render_crop,
)


def run_agent_pipeline(
    pdf_path: str,
    out_dir: Optional[str] = None,
    on_stage: Optional[Callable[[str], None]] = None,
    project_input: Optional[Dict] = None,
    source_name: Optional[str] = None,
    require_review: bool = True,
) -> Dict:
    """Run all scoped specialist waves and return a Classic-compatible report."""
    if not os.path.isfile(pdf_path):
        raise FileNotFoundError(pdf_path)

    # Keep the llm module's counters job-local (matches Classic): the job
    # log's failure-path cost estimate in jobs.py reads llm.get_cost().
    reset_cost()

    agent_dir = os.path.join(out_dir, "agent") if out_dir else None
    memory = ProjectMemory(artifact_dir=agent_dir)
    orchestrator = Orchestrator(memory=memory, on_stage=on_stage)
    usage = AgentUsage()
    orchestrator.initialize()

    orchestrator.stage("Agent ingest: PDF -> images")
    pages = convert_pdf_to_images(pdf_path)
    page_to_b64 = {page["page_number"]: page["base64"] for page in pages}
    text_dir = os.path.join(agent_dir, "text") if agent_dir else None
    page_words = attach_text_layers(pdf_path, pages, text_dir=text_dir)
    page_to_text = {page["page_number"]: page.get("text_layer") for page in pages}

    orchestrator.stage("Agent wave 1: extract sheets")
    extract_scopes = [
        AgentScope(
            scope_id=f"sheet:{page['page_number']}",
            payload={"page": page},
        )
        for page in pages
    ]
    extract_results = orchestrator.run_scopes(
        SheetExtractorAgent(usage), extract_scopes, config.EXTRACT_CONCURRENCY
    )
    sheets = [
        artifact
        for result in extract_results
        for artifact in result.artifacts
    ]
    sheets.sort(key=lambda sheet: sheet.get("page_number") or 0)
    memory.replace("sheets", sheets)
    memory.dump("01-extract.json")
    # Coverage signal: text layer present but extraction failed/empty reuses
    # the failed-scopes gap-finding path (finding built below wave 6).
    for gap_page in coverage_gaps(pages, sheets):
        orchestrator.stats.failed_scopes.append(
            f"sheet_extractor:sheet:{gap_page}: extraction gap "
            f"(text layer present, no objects extracted)"
        )

    cover_meta = derive_project_meta_from_cover(
        sheets, source_name or os.path.basename(pdf_path)
    )
    merged_input = {**cover_meta, **(project_input or {})}
    orchestrator.stage("Agent wave 2: sheet index and jurisdiction")
    index_results = orchestrator.run_scopes(
        SheetIndexAgent(usage),
        [AgentScope("sheet-index", {"sheets": sheets})],
        1,
    )
    jurisdiction_results = orchestrator.run_scopes(
        JurisdictionAgent(usage),
        [AgentScope("jurisdiction", {"project_input": merged_input})],
        1,
    )
    sheet_index = (
        index_results[0].artifacts[0]
        if index_results and index_results[0].artifacts else {}
    )
    jurisdiction = (
        jurisdiction_results[0].artifacts[0]
        if jurisdiction_results and jurisdiction_results[0].artifacts else {}
    )
    memory.replace("sheet_index", sheet_index)
    memory.replace("jurisdiction", jurisdiction)
    memory.dump("02-orient.json")

    orchestrator.stage("Agent wave 3: scoped semantic linking")
    link_results = orchestrator.run_scopes(
        LinkerAgent(usage),
        build_link_scopes(sheets),
        config.AGENT_LINK_CONCURRENCY,
    )
    clusters = [
        artifact
        for result in link_results
        for artifact in result.artifacts
    ][:config.CLUSTER_MAX]
    object_graph = build_object_graph(clusters)
    disputed_count = annotate_clusters(clusters)
    if disputed_count:
        orchestrator.stage(f"[Link] {disputed_count} clusters carry disputed extracted values")
    memory.replace("clusters", clusters)
    memory.replace("object_graph", object_graph)
    memory.dump("03-link.json")

    orchestrator.stage("Agent wave 4: per-cluster conflict critics")
    conflict_scopes = [
        AgentScope(
            scope_id=f"conflict:{cluster.get('key')}",
            payload={"cluster": cluster, "page_to_b64": page_to_b64},
        )
        for cluster in clusters
    ]
    conflict_results = orchestrator.run_scopes(
        ConflictCriticAgent(usage),
        conflict_scopes,
        config.AGENT_CONFLICT_CONCURRENCY,
    )
    conflict_findings = [
        artifact
        for result in conflict_results
        for artifact in result.artifacts
    ]
    memory.extend("findings", conflict_findings)

    orchestrator.stage("Agent wave 5: scoped specialists")
    code_results = orchestrator.run_scopes(
        CodeAgent(usage),
        build_code_scopes(sheets, jurisdiction, sheet_index),
        config.AGENT_SPECIALIST_CONCURRENCY,
    )
    construct_results = orchestrator.run_scopes(
        ConstructabilityAgent(usage),
        build_construct_scopes(clusters, conflict_findings),
        config.AGENT_SPECIALIST_CONCURRENCY,
    )
    completeness_scope = AgentScope("completeness", {
        "sheet_index": sheet_index,
        "sheet_summaries": build_sheet_summaries(sheets),
        "cluster_summary": {
            "count": len(clusters),
            "by_kind": _counts(clusters, "kind"),
        },
    })
    completeness_results = orchestrator.run_scopes(
        CompletenessAgent(usage), [completeness_scope], 1
    )
    specialist_findings = [
        artifact
        for result in code_results + construct_results + completeness_results
        for artifact in result.artifacts
    ]

    orchestrator.stage("Agent wave 5b: evidence verification")
    sheet_to_page = {str(s.get("sheet_number")): s.get("page_number")
                     for s in sheets}
    verify_targets = select_findings(
        specialist_findings, clusters,
        max_checks=config.AGENT_VERIFY_MAX_CHECKS,
        severities=config.AGENT_VERIFY_SEVERITIES,
    )
    target_indexes = {id(f): i for i, f in enumerate(specialist_findings)}
    verify_scopes = []
    for finding in verify_targets:
        cited_pages = [
            sheet_to_page[str(name)]
            for name in (finding.get("sheets") or [])
            if sheet_to_page.get(str(name)) in page_to_b64
        ]
        images = [
            page_to_b64[p]
            for p in cited_pages[:config.AGENT_CONFLICT_MAX_IMAGES]
        ]
        if not images:
            continue  # never judge evidence against images we could not load
        # Text oracle: concatenated text layer of the cited sheets, capped.
        excerpt = "\n\n".join(
            f"--- Page {p} ---\n{page_to_text[p]}"
            for p in cited_pages
            if page_to_text.get(p)
        )[:config.VERIFY_TEXT_MAX_CHARS]
        if config.VERIFY_HI_DPI_CROPS:
            images = _evidence_crops(finding, cited_pages, sheet_to_page,
                                     page_words, page_to_b64, pdf_path,
                                     fallback=images)
        verify_scopes.append(AgentScope(
            scope_id=f"verify:{target_indexes[id(finding)]}",
            payload={
                "finding_index": target_indexes[id(finding)],
                "finding": finding,
                "images_b64": images,
                "text_layer_excerpt": excerpt,
            },
        ))
    verify_results = orchestrator.run_scopes(
        EvidenceVerifierAgent(usage), verify_scopes, config.AGENT_VERIFY_CONCURRENCY)
    suppressed = apply_verdicts(specialist_findings, verify_results)
    if suppressed:
        suppressed_ids = {id(f) for f in suppressed}
        specialist_findings = [f for f in specialist_findings if id(f) not in suppressed_ids]
        memory.replace("suppressed", suppressed)

    memory.extend("findings", specialist_findings)
    gap_findings = [
        {
            "issue_id": f"AGENT-GAP-{index + 1:03d}",
            "source_stage": "qaqc",
            "category": "analysis_gap",
            "severity": "low",
            "confidence": "high",
            "location": failed_scope.split(":", 2)[1] if ":" in failed_scope else "",
            "disciplines": [],
            "sheets": [],
            "description": f"Agent analysis scope did not complete: {failed_scope}",
            "evidence": [],
            "recommended_resolution": "Review this scope manually or rerun the job.",
            "code_reference": None,
            "agent": "completeness",
            "scope_id": "failed-scopes",
        }
        for index, failed_scope in enumerate(orchestrator.stats.failed_scopes)
    ]
    memory.extend("findings", gap_findings)
    memory.dump("05-specialists.json")

    orchestrator.stage("Agent wave 6: Brain merge, judge, prioritize")
    all_findings = memory.snapshot()["findings"]
    if all_findings:
        prioritized, decisions = BrainAgent(usage).run(
            all_findings, sheet_index, jurisdiction
        )
    else:
        prioritized, decisions = [], []
    memory.extend("decisions", decisions)
    orchestrator.stats.merges = sum(
        1 for decision in decisions if decision.get("action") == "merged"
    )

    if require_review:
        orchestrator.stage("Agent review gate: build human-review queue")
        memory_snapshot = memory.snapshot()
        queue = build_review_queue(memory_snapshot, prioritized, decisions,
                                   limit=config.AGENT_REVIEW_AUDIT_SAMPLE)
        store = ReviewStore(out_dir)
        store.write_queue(queue)
        candidate_conflicts = [_finding_as_conflict(item) for item in conflict_findings]
        report = build_report(
            conflicts=candidate_conflicts,
            sheets=sheets,
            clusters=clusters,
            source=source_name or os.path.basename(pdf_path),
        )
        report.update({
            "project_input": merged_input,
            "jurisdiction": jurisdiction,
            "sheet_index": sheet_index,
            "project_intelligence": object_graph,
            "validated_issues": prioritized,
            "rfis": [],
            "suppressed_issues": memory.snapshot().get("suppressed") or [],
        })
        progress = store.progress(queue)
        # Same usage/stats summary block as the wave-7 path (rfis: 0 — they
        # are drafted only after human review finalizes the run).
        cost = usage.snapshot()
        orchestrator.stats.calls = cost["calls"]
        stats = orchestrator.stats.as_dict()
        report["summary"].update({
            "pipeline_mode": "agent",
            "agent_status": "needs_review",
            "review": progress,
            "agent_stats": stats,
            "by_stage": {
                "conflicts": len(conflict_findings),
                "qaqc": sum(
                    1 for item in specialist_findings
                    if item.get("source_stage") == "qaqc"
                ),
                "code": sum(
                    1 for item in specialist_findings
                    if item.get("source_stage") == "code"
                ),
                "constructability": sum(
                    1 for item in specialist_findings
                    if item.get("source_stage") == "constructability"
                ),
                "validated": len(prioritized),
                "rfis": 0,
            },
            "cost_usd": round(cost["usd"], 4),
            "llm_calls": cost["calls"],
            "cached_calls": cost["cached"],
            "cost_by_stage": cost["by_stage"],
            "text_backend": "openrouter",
            "models_used": cost["models"],
        })
        if out_dir:
            _dump(out_dir, "conflicts.json", report)
            _dump(out_dir, "validated_issues.json", prioritized)
        # Snapshot for the review finalizer's targeted clarification reruns.
        memory.dump("memory.json")
        return report

    orchestrator.stage("Agent wave 7: per-finding RFI writers")
    rfi_scopes = [
        AgentScope(
            scope_id=f"rfi:{finding.get('issue_id') or index + 1}",
            payload={"finding": finding},
        )
        for index, finding in enumerate(prioritized)
    ]
    rfi_results = orchestrator.run_scopes(
        RFIWriterAgent(usage), rfi_scopes, config.AGENT_RFI_CONCURRENCY
    )
    rfis = [
        artifact for result in rfi_results for artifact in result.artifacts
    ]
    memory.extend("rfis", rfis)
    memory.dump("memory.json")

    orchestrator.stage("Build agent report")
    conflicts = [_finding_as_conflict(item) for item in conflict_findings]
    report = build_report(
        conflicts=conflicts,
        sheets=sheets,
        clusters=clusters,
        source=source_name or os.path.basename(pdf_path),
    )
    report.update({
        "project_input": merged_input,
        "jurisdiction": jurisdiction,
        "sheet_index": sheet_index,
        "project_intelligence": object_graph,
        "validated_issues": prioritized,
        "rfis": rfis,
        "suppressed_issues": memory.snapshot().get("suppressed") or [],
    })
    cost = usage.snapshot()
    orchestrator.stats.calls = cost["calls"]
    stats = orchestrator.stats.as_dict()
    report["summary"].update({
        "pipeline_mode": "agent",
        "agent_status": "complete",
        "agent_stats": stats,
        "by_stage": {
            "conflicts": len(conflict_findings),
            "qaqc": sum(
                1 for item in specialist_findings
                if item.get("source_stage") == "qaqc"
            ),
            "code": sum(
                1 for item in specialist_findings
                if item.get("source_stage") == "code"
            ),
            "constructability": sum(
                1 for item in specialist_findings
                if item.get("source_stage") == "constructability"
            ),
            "validated": len(prioritized),
            "rfis": len(rfis),
        },
        "cost_usd": round(cost["usd"], 4),
        "llm_calls": cost["calls"],
        "cached_calls": cost["cached"],
        "cost_by_stage": cost["by_stage"],
        "text_backend": "openrouter",
        "models_used": cost["models"],
    })

    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
        _dump(out_dir, "assertions.json", sheets)
        _dump(out_dir, "clusters.json", [_without_base64(item) for item in clusters])
        _dump(out_dir, "sheet_index.json", sheet_index)
        _dump(out_dir, "jurisdiction.json", jurisdiction)
        _dump(out_dir, "project_intelligence.json", object_graph)
        _dump(out_dir, "validated_issues.json", prioritized)
        _dump(out_dir, "rfis.json", rfis)
        _dump(out_dir, "conflicts.json", report)
        with open(os.path.join(out_dir, "report.md"), "w", encoding="utf-8") as f:
            f.write(to_markdown(report))

    return report


def _dump(out_dir: str, name: str, value) -> None:
    with open(os.path.join(out_dir, name), "w", encoding="utf-8") as f:
        json.dump(value, f, indent=2)


def _evidence_crops(
    finding: Dict,
    cited_pages: list,
    sheet_to_page: Dict,
    page_words: Dict,
    page_to_b64: Dict,
    pdf_path: str,
    fallback: list,
) -> list:
    """High-DPI crops around each evidence item's source_text, located via the
    page text layer. Crops REPLACE full-page images when at least one evidence
    location resolves confidently; otherwise the full-page fallback is kept.
    Never returns an empty list when fallback is non-empty (I2 guard)."""
    crops: list = []
    for item in finding.get("evidence") or []:
        if len(crops) >= config.AGENT_CONFLICT_MAX_IMAGES:
            break
        if not isinstance(item, dict):
            continue
        source_text = item.get("source_text") or ""
        if not source_text:
            continue
        # Prefer the page named on the evidence item, then any cited page.
        candidates = []
        named_page = sheet_to_page.get(str(item.get("sheet") or ""))
        if named_page in cited_pages:
            candidates.append(named_page)
        candidates.extend(p for p in cited_pages if p not in candidates)
        for page in candidates:
            bbox = find_evidence_bbox(page_words.get(page) or [], source_text)
            if bbox is None:
                continue
            crop = render_crop(pdf_path, page, bbox)
            if not crop:
                continue
            crops.append(base64.b64encode(crop).decode("utf-8"))
            break
    return crops or fallback


def _counts(items, key: str) -> Dict[str, int]:
    counts: Dict[str, int] = {}
    for item in items:
        value = str(item.get(key) or "unknown")
        counts[value] = counts.get(value, 0) + 1
    return counts


def _finding_as_conflict(finding: Dict) -> Dict:
    return {
        "category": finding.get("category") or "uncategorized",
        "severity": finding.get("severity") or "medium",
        "disciplines": finding.get("disciplines") or [],
        "location": finding.get("location") or "",
        "sheets": finding.get("sheets") or [],
        "description": finding.get("description") or "",
        "evidence": finding.get("evidence") or [],
        "recommended_resolution": finding.get("recommended_resolution") or "",
        "confidence": finding.get("confidence") or "medium",
        "cluster_key": finding.get("scope_id"),
    }


def _without_base64(cluster: Dict) -> Dict:
    return {
        **cluster,
        "assertions": [
            {key: value for key, value in assertion.items() if key != "base64"}
            for assertion in cluster.get("assertions") or []
        ],
    }