Tool_Agent/tools/local/flux/bfl_client.py

"""BFL FLUX HTTP 客户端 — 异步提交 + 轮询。

文档: https://docs.bfl.ai/quick_start/generating_images
"""

from __future__ import annotations

import os
from typing import Any

import httpx
from dotenv import load_dotenv

_ = load_dotenv()

DEFAULT_API_BASE = "https://api.bfl.ai/v1"


def _api_key() -> str:
    key = os.environ.get("BFL_API_KEY", "").strip()
    if not key:
        raise ValueError("缺少环境变量 BFL_API_KEY")
    return key


def _headers() -> dict[str, str]:
    return {
        "accept": "application/json",
        "x-key": _api_key(),
        "Content-Type": "application/json",
    }


def submit_generation(
    *,
    model: str,
    prompt: str,
    width: int | None = None,
    height: int | None = None,
    parameters: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """POST {BFL_API_BASE}/{model}，返回含 id、polling_url 等（以 BFL 响应为准）。"""
    base = os.environ.get("BFL_API_BASE", DEFAULT_API_BASE).rstrip("/")
    model_path = model.strip().lstrip("/")
    url = f"{base}/{model_path}"

    body: dict[str, Any] = dict(parameters) if parameters else {}
    body["prompt"] = prompt
    if width is not None:
        body["width"] = width
    if height is not None:
        body["height"] = height

    with httpx.Client(timeout=120.0) as client:
        r = client.post(url, headers=_headers(), json=body)
        try:
            data = r.json()
        except Exception:
            r.raise_for_status()
            raise RuntimeError(r.text[:2000]) from None

    if r.status_code >= 400:
        err = data.get("detail") if isinstance(data, dict) else None
        msg = err if err is not None else str(data)
        raise RuntimeError(f"BFL HTTP {r.status_code}: {msg}")

    if not isinstance(data, dict):
        raise RuntimeError("提交响应不是 JSON 对象")
    return data


def poll_result(*, polling_url: str, request_id: str) -> dict[str, Any]:
    """GET polling_url，Query: id=request_id（与官方示例一致）。"""
    with httpx.Client(timeout=60.0) as client:
        r = client.get(
            polling_url.strip(),
            headers={
                "accept": "application/json",
                "x-key": _api_key(),
            },
            params={"id": request_id.strip()},
        )
        try:
            data = r.json()
        except Exception:
            r.raise_for_status()
            raise RuntimeError(r.text[:2000]) from None

    if r.status_code >= 400:
        msg = data if isinstance(data, dict) else str(data)
        raise RuntimeError(f"BFL poll HTTP {r.status_code}: {msg}")

    if not isinstance(data, dict):
        raise RuntimeError("轮询响应不是 JSON 对象")
    return data