scripts/test_aggregator.py

"""End-to-end smoke test: Image Analyzer -> Planner -> Vision Pass -> Aggregator.

Runs the full pipeline so far on the sample images and prints the
aggregated findings + any conflicts that the review agent would address.

Usage:
    uv run python scripts/test_aggregator.py
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

from dotenv import load_dotenv

from ergo_agentic.datasources import DatasourceRegistry
from ergo_agentic.models import DEFAULT_MODEL_CONFIG
from ergo_agentic.nodes.aggregator import aggregate
from ergo_agentic.nodes.image_analyzer import analyze_image
from ergo_agentic.nodes.parameter_planner import plan_parameters
from ergo_agentic.nodes.vision_pass import _make_node
from ergo_agentic.state import ImageInput

REPORT_SAMPLE_FILE = (
    Path(__file__).resolve().parents[1] / "docs" / "datasources" / "report-sample.json"
)


def main() -> int:
    load_dotenv()
    registry = DatasourceRegistry.from_knowledge_base()
    run_vision_pass = _make_node(registry)

    with REPORT_SAMPLE_FILE.open() as f:
        report = json.load(f)
    urls = report.get("uploadedImages", [])
    images: list[ImageInput] = [
        {"image_id": f"img_{i}", "url": url, "label": None}
        for i, url in enumerate(urls, start=1)
    ]

    print(f"Step 1: analyzing {len(images)} images...")
    manifests = []
    for img in images:
        result = analyze_image(
            {"image": img, "model_id": DEFAULT_MODEL_CONFIG.image_analyzer}
        )
        manifests.append(result["image_manifests"][0])

    print("Step 2: planning parameters...")
    plan_state = {"image_manifests": manifests}
    plan_result = plan_parameters(plan_state, registry=registry)

    print("Step 3: running vision passes...")
    images_by_id = {img["image_id"]: img for img in images}
    model_id = DEFAULT_MODEL_CONFIG.vision_passes[0]
    observations = []
    for fg_name, fg_plan in plan_result["execution_plan"]["focus_groups"].items():
        if not fg_plan["parameter_ids"] or not fg_plan["image_ids"]:
            continue
        for image_id in fg_plan["image_ids"]:
            result = run_vision_pass({
                "image": images_by_id[image_id],
                "model_id": model_id,
                "focus_group": fg_name,
                "parameter_ids": fg_plan["parameter_ids"],
            })
            observations.extend(result["observations"])

    print(f"  collected {len(observations)} observations")

    print("\nStep 4: aggregating...\n")
    agg_state = {"observations": observations}
    agg_result = aggregate(agg_state, registry=registry)

    findings = agg_result["aggregated_findings"]
    conflicts = agg_result["conflicts"]

    for pid, finding in findings.items():
        param = registry.get_parameter(pid)
        label = param.parameter_text if param else pid
        print(f"[{label}] selection_mode={finding.selection_mode.value}")
        if not finding.worst_outcomes:
            print("  (no visible observations — skipped)")
            continue
        print(f"  candidates: {finding.candidate_outcomes}")
        print(f"  worst:      {finding.worst_outcomes}")
        if finding.has_conflict:
            print(f"  CONFLICT: {finding.conflict_detail}")
        for outcome, evidence in finding.evidence.items():
            opt = registry.get_option(outcome)
            risk = opt.risk_level.value if opt and opt.risk_level else "good"
            score = opt.posture_score if opt else "?"
            print(
                f"    -> {outcome} [risk={risk}, score={score}, "
                f"agreement={evidence.agreement_level.value}, "
                f"images={evidence.source_images}]"
            )
        print()

    print(f"Total findings: {len(findings)}, conflicts: {len(conflicts)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())