283 lines
11 KiB
Python
283 lines
11 KiB
Python
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
|
|
|
|
|
|
DIAGNOSTIC_SCALAR_KEYS = {
|
|
"code",
|
|
"error",
|
|
"error_code",
|
|
"error_message",
|
|
"id",
|
|
"is_final",
|
|
"message",
|
|
"msg",
|
|
"state",
|
|
"status",
|
|
"status_group",
|
|
"task_id",
|
|
"type",
|
|
}
|
|
|
|
|
|
def safe_response_diagnostic(payload: Any) -> dict[str, Any]:
|
|
"""Describe provider payload shape without persisting images, URLs, or secrets."""
|
|
|
|
signals: dict[str, Any] = {}
|
|
|
|
def shape(value: Any, path: str = "root", depth: int = 0) -> Any:
|
|
if depth > 4:
|
|
return type(value).__name__
|
|
if isinstance(value, dict):
|
|
result: dict[str, Any] = {}
|
|
for key, child in list(value.items())[:40]:
|
|
child_path = f"{path}.{key}"
|
|
if key in DIAGNOSTIC_SCALAR_KEYS and isinstance(child, (str, int, float, bool)):
|
|
rendered = str(child)
|
|
signals[child_path] = rendered[:500]
|
|
if key in {"url", "result_url", "b64_json", "data", "inlineData", "inline_data"}:
|
|
if isinstance(child, list):
|
|
result[key] = {"type": "list", "length": len(child)}
|
|
elif isinstance(child, dict):
|
|
result[key] = {"type": "object", "keys": list(child)[:30]}
|
|
elif isinstance(child, str):
|
|
result[key] = {"type": "string", "length": len(child)}
|
|
else:
|
|
result[key] = type(child).__name__
|
|
if isinstance(child, (dict, list)):
|
|
nested = shape(child, child_path, depth + 1)
|
|
result[f"{key}_shape"] = nested
|
|
continue
|
|
result[key] = shape(child, child_path, depth + 1)
|
|
return result
|
|
if isinstance(value, list):
|
|
return {
|
|
"type": "list",
|
|
"length": len(value),
|
|
"items": [shape(item, f"{path}[{index}]", depth + 1) for index, item in enumerate(value[:3])],
|
|
}
|
|
if value is None:
|
|
return "null"
|
|
if isinstance(value, str):
|
|
return {"type": "string", "length": len(value)}
|
|
return type(value).__name__
|
|
|
|
return {"shape": shape(payload), "signals": signals}
|
|
|
|
|
|
def extract_task_id(payload: Any) -> str | int | None:
|
|
if not isinstance(payload, dict):
|
|
return None
|
|
direct = payload.get("task_id")
|
|
if direct not in (None, ""):
|
|
return direct
|
|
data = payload.get("data")
|
|
if isinstance(data, dict):
|
|
nested = data.get("task_id")
|
|
if nested not in (None, ""):
|
|
return nested
|
|
for key in ("task_ids", "任务ids"):
|
|
candidates = data.get(key)
|
|
if isinstance(candidates, list) and candidates:
|
|
return candidates[0]
|
|
return None
|
|
|
|
|
|
def unwrap_media_status(payload: Any) -> dict[str, Any]:
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError("AIGC 状态接口返回的不是 JSON 对象。")
|
|
data = payload.get("data")
|
|
if isinstance(data, dict) and any(
|
|
key in data for key in ("state", "is_final", "result_url", "error", "progress")
|
|
):
|
|
return {**payload, **data}
|
|
return payload
|
|
|
|
|
|
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()
|
|
created = response.json()
|
|
task_id = extract_task_id(created)
|
|
if task_id is not None:
|
|
return await self._poll_aigc_media(client, item, task_id, created)
|
|
return created
|
|
|
|
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 = extract_task_id(created)
|
|
if task_id is None:
|
|
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 = unwrap_media_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
|