64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Resume one known API Mart image task and archive its traceable output."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
|
|
DEFAULT_WRAPPER = Path.home() / ".agents/skills/image-gen/scripts/apimart_curl.zsh"
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--task-id", required=True)
|
|
parser.add_argument("--output", required=True, type=Path)
|
|
parser.add_argument("--report", required=True, type=Path)
|
|
parser.add_argument("--wrapper", type=Path, default=DEFAULT_WRAPPER)
|
|
args = parser.parse_args()
|
|
|
|
response = subprocess.run(
|
|
[str(args.wrapper), f"https://api.apimart.ai/v1/tasks/{args.task_id}"],
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
payload = json.loads(response.stdout)
|
|
data = payload.get("data", {})
|
|
status = str(data.get("status", "")).lower()
|
|
if status != "completed":
|
|
raise SystemExit(f"task {args.task_id} is {status or 'unknown'}, not completed")
|
|
images = data.get("result", {}).get("images", [])
|
|
urls = images[0].get("url", []) if images else []
|
|
if not urls:
|
|
raise SystemExit(f"task {args.task_id} has no output URL")
|
|
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
with urllib.request.urlopen(str(urls[0]), timeout=120) as source:
|
|
args.output.write_bytes(source.read())
|
|
report = {
|
|
"schema_version": 1,
|
|
"provider": "API Mart",
|
|
"model": "gpt-image-2",
|
|
"task_id": args.task_id,
|
|
"status": status,
|
|
"cost_usd": data.get("cost"),
|
|
"credits_cost": data.get("credits_cost"),
|
|
"created": data.get("created"),
|
|
"completed": data.get("completed"),
|
|
"actual_time_seconds": data.get("actual_time"),
|
|
"bytes": args.output.stat().st_size,
|
|
"output": str(args.output),
|
|
}
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
print(f"ABG_APIMART_IMAGE_TASK_OK task={args.task_id} output={args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|