ai-image/app/providers/base.py

"""供应商适配器基类与能力声明。"""

from __future__ import annotations

import time
from dataclasses import dataclass

import httpx


@dataclass
class ProviderConfig:
    """一条供应商配置（api_key 已解密）。"""

    id: str
    name: str
    adapter: str
    base_url: str
    api_key: str
    default_model: str = ""


@dataclass
class Capabilities:
    """适配器能力声明，服务端据此校验入参。"""

    supports_image_input: bool = False
    size_hint: str = "WxH"
    kind: str = "image"


class ProviderError(Exception):
    """上游调用失败。"""

    def __init__(self, message: str, status: int = 502):
        super().__init__(message)
        self.message = message
        self.status = status


@dataclass
class ImageResult:
    """单张出图结果：要么有二进制数据，要么有远端 URL。"""

    data: bytes | None = None
    url: str | None = None
    filename: str | None = None
    content_type: str = "image/png"


class ProviderAdapter:
    """适配器基类。子类需实现 list_models 与 generate。"""

    name = "base"
    capabilities = Capabilities()

    def __init__(self, cfg: ProviderConfig):
        self.cfg = cfg

    # ----------------------------------------------------------- 工具方法
    def _headers(self) -> dict[str, str]:
        headers = {"Accept": "application/json"}
        if self.cfg.api_key:
            headers["Authorization"] = f"Bearer {self.cfg.api_key}"
        return headers

    def _url(self, path: str) -> str:
        return self.cfg.base_url.rstrip("/") + "/" + path.lstrip("/")

    def _request(
        self,
        method: str,
        path: str,
        *,
        json_body: dict | None = None,
        timeout: float = 60.0,
        retries: int = 2,
    ) -> httpx.Response:
        url = self._url(path)
        last_exc: Exception | None = None
        for attempt in range(retries + 1):
            try:
                with httpx.Client(timeout=timeout, follow_redirects=True) as client:
                    resp = client.request(method, url, headers=self._headers(), json=json_body)
            except httpx.HTTPError as exc:
                last_exc = exc
            else:
                if resp.status_code < 400:
                    return resp
                if resp.status_code in (429, 500, 502, 503, 504) and attempt < retries:
                    time.sleep(1.5 * (attempt + 1))
                    continue
                raise ProviderError(
                    f"上游返回 HTTP {resp.status_code}: {resp.text[:400]}",
                    status=resp.status_code,
                )
            if attempt < retries:
                time.sleep(1.5 * (attempt + 1))
        raise ProviderError(f"请求上游失败: {last_exc}")

    def _json(self, resp: httpx.Response) -> dict:
        try:
            payload = resp.json()
        except ValueError:
            raise ProviderError(f"上游返回非 JSON 内容: {resp.text[:200]}")
        if not isinstance(payload, dict):
            raise ProviderError("上游 JSON 响应必须是对象")
        return payload

    # ----------------------------------------------------------- 接口定义
    def list_models(self) -> list[str]:
        payload = self._json(self._request("GET", "/models"))
        items = payload.get("data") or payload.get("models") or []
        if not isinstance(items, list):
            raise ProviderError("上游模型列表格式错误")
        ids: list[str] = []
        for item in items:
            mid = (item.get("id") or item.get("name")) if isinstance(item, dict) else item
            if isinstance(mid, str) and mid and mid not in ids:
                ids.append(mid)
        return ids

    def generate(
        self,
        model: str,
        prompt: str,
        size: str,
        ratio: str,
        n: int,
        images: list[str],
    ) -> list[ImageResult]:
        raise NotImplementedError

    # ----------------------------------------------------------- 公共逻辑
    @staticmethod
    def filter_image_models(model_ids: list[str]) -> list[str]:
        """命中规则：模型 ID 转小写后包含 image。"""
        return [m for m in model_ids if "image" in m.lower()]

    def download(self, url: str) -> ImageResult:
        try:
            with httpx.Client(timeout=120.0, follow_redirects=True) as client:
                resp = client.get(url)
        except httpx.HTTPError as exc:
            raise ProviderError(f"下载图片失败: {exc}")
        if resp.status_code >= 400:
            raise ProviderError(f"下载图片失败: HTTP {resp.status_code}")
        ctype = resp.headers.get("content-type", "image/png")
        return ImageResult(data=resp.content, content_type=ctype, filename=_ext_of(ctype))

    @staticmethod
    def collect_results(payload: dict) -> list[ImageResult]:
        """从 data[0].b64_json / data[0].url 两种形态中取图。"""
        import base64
        import binascii

        items = payload.get("data") or []
        if not isinstance(items, list):
            raise ProviderError("上游图片列表格式错误")
        results: list[ImageResult] = []
        for item in items:
            if not isinstance(item, dict):
                continue
            b64 = item.get("b64_json")
            if b64:
                try:
                    data = base64.b64decode(b64, validate=True)
                except (binascii.Error, ValueError, TypeError) as exc:
                    raise ProviderError("上游返回无效的 Base64 图片") from exc
                if not data:
                    raise ProviderError("上游返回空图片")
                results.append(
                    ImageResult(data=data, content_type="image/png", filename=".png")
                )
                continue
            url = item.get("url")
            if url:
                results.append(ImageResult(url=url, filename=_ext_of(url)))
        if not results:
            err = payload.get("error")
            if err:
                raise ProviderError(f"上游错误: {err}")
            raise ProviderError(f"上游未返回图片: {str(payload)[:300]}")
        return results


def _ext_of(value: str) -> str:
    v = value.lower()
    for ext in (".png", ".jpg", ".jpeg", ".webp"):
        if ext in v:
            return ext
    return ".png"