scripts/test_planner.py

"""Smoke test for the Parameter Planner.

Runs the Image Analyzer on the sample report images, then runs the planner
on the resulting manifests to show what would be assessed.

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

from __future__ import annotations

import asyncio
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.image_analyzer import analyze_image
from ergo_agentic.nodes.parameter_planner import plan_parameters
from ergo_agentic.nodes.routing import build_routing_manifest
from ergo_agentic.state import ImageInput

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


async def _main() -> int:
    load_dotenv()
    registry = DatasourceRegistry.from_knowledge_base()

    with REPORT_SAMPLE_FILE.open() as f:
        report = json.load(f)
    urls = report.get("uploadedImages", [])

    print(f"Step 1: Analyzing {len(urls)} images...")
    manifests = []
    for i, url in enumerate(urls, start=1):
        image: ImageInput = {"image_id": f"img_{i}", "url": url, "label": None}
        result = await analyze_image(
            {"image": image, "model_id": DEFAULT_MODEL_CONFIG.image_analyzer}
        )
        manifests.append(result["image_manifests"][0])
        print(f"  img_{i}: location={manifests[-1].work_location}, "
              f"body_coverage={manifests[-1].body_coverage}, "
              f"screens={len(manifests[-1].screens)}, "
              f"posture_hint={manifests[-1].posture_context_hint}")

    print("\nStep 2: Running Parameter Planner...\n")
    routing_manifest = build_routing_manifest(
        manifests=manifests,
        cv_results=[],
        metadata={},
    )
    state = {"image_manifests": manifests, "routing_manifest": routing_manifest}
    plan_result = plan_parameters(state, registry=registry)

    scene = plan_result["scene_config"]
    print("Scene config:")
    print(f"  screen_count: {scene.screen_count}")
    print(f"  screen_types: {scene.screen_types}")
    print(f"  has_standing_desk: {scene.has_standing_desk}")
    print(f"  work_location: {scene.work_location}")
    print(f"  person_detected: {scene.person_detected}")

    print("\nExecution plan:")
    for fg_name, plan in plan_result["execution_plan"]["focus_groups"].items():
        print(f"\n  [{fg_name}]")
        if plan["skip_reason"]:
            print(f"    SKIPPED: {plan['skip_reason']}")
            continue
        print(f"    images: {plan['image_ids']}")
        print(f"    parameters ({len(plan['parameter_ids'])}):")
        for pid in plan["parameter_ids"]:
            p = registry.get_parameter(pid)
            sample_key = p.options[0].key
            print(f"      - {p.parameter_text} ({sample_key.rsplit('-', 1)[0]}-*)")

    print(f"\nTotal assessable: {len(plan_result['assessable_parameters'])}")
    print(f"Total skipped: {len(plan_result['skipped_parameters'])}")
    if plan_result["skipped_parameters"]:
        print("Skipped parameters:")
        for pid in plan_result["skipped_parameters"]:
            p = registry.get_parameter(pid)
            sample_key = p.options[0].key
            print(f"  - {p.parameter_text} ({sample_key.rsplit('-', 1)[0]}-*)")

    return 0


def main() -> int:
    return asyncio.run(_main())


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