Poll an asynchronous image task safely

import base64
import os
import random
import time
from pathlib import Path

import requests

task_id = os.environ["YINGTU_TASK_ID"]
deadline = time.monotonic() + 180
delay = 2.0

while True:
    if time.monotonic() >= deadline:
        raise TimeoutError(f"Polling deadline reached for {task_id}")

    response = requests.get(
        f"https://api.yingtu.ai/v1/tasks/{task_id}",
        headers={"Authorization": f"Bearer {os.environ['YINGTU_API_KEY']}"},
        timeout=30,
    )

    if response.status_code == 429:
        retry_after = response.headers.get("Retry-After")
        wait = float(retry_after) if retry_after else delay
        time.sleep(min(wait, max(0, deadline - time.monotonic())))
        delay = min(delay * 1.5, 10)
        continue

    response.raise_for_status()
    task = response.json()
    status = task["status"]

    if status == "completed":
        model = task["model"]
        if model.startswith("gemini-"):
            image = next(
                part["inlineData"]
                for candidate in task["result"].get("candidates", [])
                for part in candidate.get("content", {}).get("parts", [])
                if part.get("inlineData", {}).get("data")
            )
            suffix = {
                "image/png": ".png",
                "image/jpeg": ".jpg",
                "image/webp": ".webp",
            }[image["mimeType"]]
            encoded = image["data"]
        else:
            image = task["result"]["data"][0]
            encoded = image["b64_json"]
            suffix = ".png"  # documented reverse-image routes return PNG

        Path(f"{task_id}{suffix}").write_bytes(base64.b64decode(encoded))
        break

    if status == "failed":
        error = task.get("error", {})
        raise RuntimeError(
            f"{error.get('code', 'task_failed')}: "
            f"{error.get('message', 'unknown task failure')}"
        )

    if status == "expired":
        raise RuntimeError("result_expired: create a new task if output is still needed")

    if status not in {"queued", "in_progress"}:
        raise RuntimeError(f"Unexpected task status: {status}")

    time.sleep(delay + random.uniform(0, min(0.5, delay * 0.1)))
    delay = min(delay * 1.5, 10)