examples/decision.py

"""Send the recorded fictional scenarios to a live Premonition decision server."""

import argparse
import base64
import json
from pathlib import Path

import httpx


def request(case, root):
    media = []
    for name in case["files"]:
        path = root / "demo" / "assets" / name
        mime, kind = {
            ".png": ("image/png", "image_url"),
            ".mp4": ("video/mp4", "video_url"),
            ".wav": ("audio/wav", "audio_url"),
        }[path.suffix]
        encoded = base64.b64encode(path.read_bytes()).decode()
        media.append({"type": kind, kind: {"url": f"data:{mime};base64,{encoded}"}})
    return {
        "model": "Premonition-Exp",
        "state": case["state"],
        "questions": case["questions"],
        "media": media,
        "samples": 1,
        "seed": 42,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="http://127.0.0.1:8011")
    parser.add_argument("--case", default="omni-review")
    parser.add_argument("--all", action="store_true")
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    args = parser.parse_args()
    rows = json.loads((args.root / "demo/cases.json").read_text())
    selected = rows if args.all else [c for c in rows if c["id"] == args.case]
    if not selected:
        parser.error("Unknown case. Choose from: " + ", ".join(c["id"] for c in rows))
    with httpx.Client(timeout=300) as client:
        for case in selected:
            response = client.post(
                args.url.rstrip("/") + "/v1/systemone", json=request(case, args.root)
            )
            response.raise_for_status()
            print(
                json.dumps({"case": case["id"], "response": response.json()}, indent=2), flush=True
            )


if __name__ == "__main__":
    main()