scripts/test_vision_pass.py

"""Smoke test for the focused Vision Pass.

Runs the full pipeline so far: Image Analyzer → Parameter Planner → Vision
Passes (sequential for clarity, not parallel). Prints the observations
produced for each (image, model, focus_group) combination.

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

from __future__ import annotations

import json
import sys
from pathlib import Path

from dotenv import load_dotenv

from ergo_agentic.datasources import DatasourceRegistry

REPORT_SAMPLE_FILE = (
    Path(__file__).resolve().parents[1] / "docs" / "datasources" / "report-sample.json"
)
from ergo_agentic.domain.enums import FocusGroup
from ergo_agentic.models import DEFAULT_MODEL_CONFIG
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


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("\nStep 2: Planning parameters...")
    plan_result = plan_parameters(
        {"image_manifests": manifests}, registry=registry
    )
    plan = plan_result["execution_plan"]

    images_by_id = {img["image_id"]: img for img in images}
    model_id = DEFAULT_MODEL_CONFIG.vision_passes[0]

    print(f"\nStep 3: Running vision passes (model: {model_id})...\n")
    all_observations = []

    for fg_name, fg_plan in plan["focus_groups"].items():
        if not fg_plan["parameter_ids"] or not fg_plan["image_ids"]:
            continue
        print(f"\n{'='*60}")
        print(f"Focus group: {fg_name}")
        print(f"{'='*60}")
        for image_id in fg_plan["image_ids"]:
            img = images_by_id[image_id]
            print(f"\n--- {image_id} ---")
            try:
                result = run_vision_pass({
                    "image": img,
                    "model_id": model_id,
                    "focus_group": fg_name,
                    "parameter_ids": fg_plan["parameter_ids"],
                })
                for obs in result["observations"]:
                    p = registry.get_parameter(obs.parameter_id)
                    label = p.parameter_text if p else obs.parameter_id
                    outcomes = ", ".join(obs.selected_outcomes) or "(none)"
                    print(f"  [{label}]")
                    print(f"    visibility: {obs.visibility.value} | confidence: {obs.confidence.value}")
                    print(f"    selected: {outcomes}")
                    print(f"    note: {obs.evidence_note}")
                    all_observations.append(obs)
            except Exception as e:
                print(f"  ERROR: {e}")

    print(f"\n\nTotal observations: {len(all_observations)}")
    return 0


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