"""供应商适配器基类与能力声明。""" 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"