import asyncio import base64 from collections.abc import Mapping from typing import Any from urllib.parse import urljoin import httpx from app.integrations.model_router import ModelRouter class OpenAICompatibleGateway: """Adapter for independently configured OpenAI-compatible model-pool items.""" def __init__(self, values: Mapping[str, Any]) -> None: self.values = values self.router = ModelRouter(values) @staticmethod def _url(item: Mapping[str, Any], path_key: str, fallback: str) -> str: base_url = str(item.get("base_url", "")).rstrip("/") + "/" path = str(item.get(path_key, fallback)).lstrip("/") return urljoin(base_url, path) @staticmethod def _headers(item: Mapping[str, Any], *, json_content: bool = True) -> dict[str, str]: headers = {"Authorization": f"Bearer {item.get('api_key', '')}"} if json_content: headers["Content-Type"] = "application/json" return headers async def list_models(self, model: Mapping[str, Any] | None = None) -> list[str]: item = model or self.router.orchestrator() async with httpx.AsyncClient(timeout=15) as client: response = await client.get( self._url(item, "models_path", "/models"), headers=self._headers(item), ) response.raise_for_status() payload = response.json() return [candidate["id"] for candidate in payload.get("data", []) if isinstance(candidate, dict) and candidate.get("id")] async def chat( self, messages: list[dict[str, Any]], *, vision: bool = False, model_instance_id: str | None = None, ) -> dict[str, Any]: item = self.router.choose("spatial", model_instance_id) if vision else self.router.orchestrator() async with httpx.AsyncClient(timeout=120) as client: response = await client.post( self._url(item, "chat_path", "/chat/completions"), headers=self._headers(item), json={"model": item["model_id"], "messages": messages}, ) response.raise_for_status() return response.json() async def generate_image( self, prompt: str, *, model_instance_id: str | None = None, **options: Any, ) -> dict[str, Any]: item = self.router.choose("image", model_instance_id) if item.get("image_protocol", "openai_images") == "aigc_media": return await self._generate_aigc_media(item, prompt, options) payload = {"model": item["model_id"], "prompt": prompt, **options} async with httpx.AsyncClient(timeout=180) as client: response = await client.post( self._url(item, "image_generation_path", "/images/generations"), headers=self._headers(item), json=payload, ) response.raise_for_status() return response.json() async def edit_image( self, prompt: str, image: bytes, *, filename: str = "image.png", mask: bytes | None = None, model_instance_id: str | None = None, **options: Any, ) -> dict[str, Any]: item = self.router.choose("image", model_instance_id) if item.get("image_protocol", "openai_images") == "aigc_media": encoded = base64.b64encode(image).decode("ascii") references = [f"data:image/png;base64,{encoded}"] if mask is not None: mask_encoded = base64.b64encode(mask).decode("ascii") references.append(f"data:image/png;base64,{mask_encoded}") options["images"] = [*references, *list(options.pop("images", []))] return await self._generate_aigc_media(item, prompt, options) files: dict[str, tuple[str, bytes, str]] = {"image": (filename, image, "image/png")} if mask is not None: files["mask"] = ("mask.png", mask, "image/png") data = {"model": item["model_id"], "prompt": prompt, **options} async with httpx.AsyncClient(timeout=180) as client: response = await client.post( self._url(item, "image_edit_path", "/images/edits"), headers=self._headers(item, json_content=False), data=data, files=files, ) response.raise_for_status() return response.json() async def _generate_aigc_media( self, item: Mapping[str, Any], prompt: str, options: Mapping[str, Any], ) -> dict[str, Any]: params, notify_url = self._media_params(item, options) payload: dict[str, Any] = {"model": item["model_id"], "prompt": prompt, "params": params} if notify_url: payload["notify_url"] = notify_url async with httpx.AsyncClient(timeout=180) as client: response = await client.post( self._url(item, "image_generation_path", "/v1/media/generate"), headers=self._headers(item), json=payload, ) response.raise_for_status() created = response.json() task_id = created.get("task_id") if isinstance(created, dict) else None if not task_id: return created return await self._poll_aigc_media(client, item, task_id, created) async def _poll_aigc_media( self, client: httpx.AsyncClient, item: Mapping[str, Any], task_id: str | int, created: dict[str, Any], ) -> dict[str, Any]: for _ in range(45): await asyncio.sleep(3) response = await client.get( self._url(item, "image_status_path", "/v1/media/status"), headers=self._headers(item), params={"task_id": task_id}, ) response.raise_for_status() status = response.json() if not status.get("is_final"): continue if status.get("state") != "success": reason = status.get("error") or status.get("status") or "平台返回失败状态" raise RuntimeError(f"AIGC 任务 {task_id} 失败：{reason}") result_url = status.get("result_url") if result_url: return {**created, **status, "data": [{"url": result_url}]} return {**created, **status} raise RuntimeError(f"AIGC 任务 {task_id} 在 135 秒内没有完成") @staticmethod def _media_params( item: Mapping[str, Any], options: Mapping[str, Any], ) -> tuple[dict[str, Any], str]: raw = {key: value for key, value in options.items() if value is not None} notify_url = str(raw.pop("notify_url", "")) profile = item.get("image_parameter_profile", "generic") if profile == "gpt_image_2": allowed = {"images", "size", "quality"} elif profile == "nano_banana_pro": if "aspect_ratio" in raw and "aspectRatio" not in raw: raw["aspectRatio"] = raw.pop("aspect_ratio") if "size" in raw and "imageSize" not in raw: raw["imageSize"] = raw.pop("size") allowed = {"images", "aspectRatio", "imageSize"} elif profile == "seedream_5_pro": if "aspectRatio" in raw and "aspect_ratio" not in raw: raw["aspect_ratio"] = raw.pop("aspectRatio") if "imageSize" in raw and "size" not in raw: raw["size"] = raw.pop("imageSize") allowed = {"images", "aspect_ratio", "size"} else: allowed = set(raw) return {key: value for key, value in raw.items() if key in allowed}, notify_url