scripts/capture_cv_samples.py

"""Capture sample responses from the CV Tool API.

Usage:
    uv run python scripts/capture_cv_samples.py path/to/image.jpg
    uv run python scripts/capture_cv_samples.py https://example.com/image.png
    uv run python scripts/capture_cv_samples.py img.jpg --heavy

Reads CV_API_BASE_URL and CV_API_KEY from the environment. Writes one JSON
per tool to tests/fixtures/cv/<tool>.json — these become the source of truth
for the typed `result` models in cv/signals.py and for snapshot tests.

Re-run when the CV service schema changes. The fixtures pin our expectations;
shape drift surfaces when the typed models fail to validate.

Heavy outputs (segmentation masks, depth maps) are omitted by default to keep
fixtures small. Pass --heavy to include them when needed for annotation work.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import sys
from pathlib import Path

from ergo_agentic.cv import HttpCVClient
from ergo_agentic.cv.client import (
    TOOL_DEPTH,
    TOOL_MEDIAPIPE_HANDS,
    TOOL_MEDIAPIPE_POSE,
    TOOL_RFDETR,
    TOOL_SAM2,
    TOOL_YOLO_DETECT,
    TOOL_YOLO_POSE,
    TOOL_YOLO_SEG,
)

# All CV signal-producing tools. gemma-ergonomics is intentionally excluded —
# it's a VLM-with-structured-output, treated as a peer model in vision_pass,
# not a CV signal source. sam2-segment-boxes is excluded because it requires
# input boxes (not a one-shot capture).
ALL_TOOLS: tuple[str, ...] = (
    TOOL_MEDIAPIPE_POSE,
    TOOL_MEDIAPIPE_HANDS,
    TOOL_YOLO_POSE,
    TOOL_YOLO_DETECT,
    TOOL_YOLO_SEG,
    TOOL_SAM2,
    TOOL_RFDETR,
    TOOL_DEPTH,
)

HEAVY_TOOLS: frozenset[str] = frozenset({TOOL_YOLO_SEG, TOOL_SAM2, TOOL_DEPTH, TOOL_RFDETR})


async def _run() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("image", help="image path or URL to send to each tool")
    parser.add_argument(
        "--base-url",
        default=os.environ.get("CV_API_BASE_URL", "http://localhost:7860"),
        help="CV service base URL (default: env CV_API_BASE_URL or http://localhost:7860)",
    )
    parser.add_argument(
        "--api-key",
        default=os.environ.get("CV_API_KEY"),
        help="API key for X-Tool-API-Key header (default: env CV_API_KEY)",
    )
    parser.add_argument(
        "--out",
        default="tests/fixtures/cv",
        help="output directory (default: tests/fixtures/cv)",
    )
    parser.add_argument(
        "--heavy",
        action="store_true",
        help="request heavy outputs (masks, depth maps). Larger fixtures.",
    )
    parser.add_argument(
        "--tools",
        nargs="+",
        choices=list(ALL_TOOLS),
        help="subset of tools to capture (default: all)",
    )
    args = parser.parse_args()

    if not args.api_key:
        parser.error("CV_API_KEY env var or --api-key required")

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    image = {"image_id": "sample", "url": args.image, "label": None}
    tools = tuple(args.tools) if args.tools else ALL_TOOLS

    print(f"CV service: {args.base_url}")
    print(f"Image:      {args.image}")
    print(f"Output:     {out_dir}")
    print(f"Heavy:      {args.heavy}")
    print()

    client = HttpCVClient(base_url=args.base_url, api_key=args.api_key)
    failures = 0
    try:
        for tool in tools:
            include_heavy = args.heavy and tool in HEAVY_TOOLS
            print(f"  → {tool:24s}", end=" ", flush=True)
            try:
                resp = await client.call(tool, image, include_heavy_outputs=include_heavy)
                out_path = out_dir / f"{tool}.json"
                out_path.write_text(json.dumps(resp.model_dump(), indent=2, default=str))
                result_size = len(json.dumps(resp.result, default=str))
                tag = " [heavy omitted]" if resp.heavy_outputs_omitted else ""
                tag += f" [warnings: {len(resp.warnings)}]" if resp.warnings else ""
                print(f"{resp.status:7s}  result={result_size:>7d}B{tag}")
            except Exception as e:  # noqa: BLE001
                print(f"FAILED  {type(e).__name__}: {e}")
                failures += 1
    finally:
        await client.aclose()

    print()
    print(f"Done. Wrote {len(tools) - failures}/{len(tools)} fixtures to {out_dir}")
    return 0 if failures == 0 else 1


if __name__ == "__main__":
    sys.exit(asyncio.run(_run()))