feat: support aggregated image model protocols

This commit is contained in:
Codex
2026-08-01 23:38:31 +08:00
parent 4430149ee4
commit 52d877836c
8 changed files with 474 additions and 9 deletions
+16
View File
@@ -12,6 +12,7 @@ from app.runtime_settings import (
SettingsReadiness,
generate_secret,
get_runtime_store,
test_image_generation,
test_model_pool_item,
test_runtime_settings,
)
@@ -79,6 +80,21 @@ async def test_model(
)
@router.post("/settings/models/{model_id}/test-generation", response_model=RuntimeSettingsTestResult)
async def test_model_generation(
model_id: str,
store: EncryptedSettingsStore = Depends(get_runtime_store),
) -> RuntimeSettingsTestResult:
try:
return await test_image_generation(store.merged_values(), model_id)
except (RuntimeError, ValueError) as exc:
return RuntimeSettingsTestResult(
target=f"generation:{model_id}",
ok=False,
message=f"真实生图失败:{exc}",
)
@router.post("/settings/generate-secret", response_model=SecretGenerateResponse)
def create_secret(request: SecretGenerateRequest) -> SecretGenerateResponse:
return SecretGenerateResponse(value=generate_secret(request.kind))
@@ -1,3 +1,5 @@
import asyncio
import base64
from collections.abc import Mapping
from typing import Any
from urllib.parse import urljoin
@@ -63,6 +65,8 @@ class OpenAICompatibleGateway:
**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(
@@ -84,6 +88,14 @@ class OpenAICompatibleGateway:
**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")
@@ -97,3 +109,79 @@ class OpenAICompatibleGateway:
)
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
+142 -2
View File
@@ -61,6 +61,8 @@ class RuntimeSettingsTestResult(BaseModel):
ModelCategory = Literal["language", "multimodal"]
ModelCapability = Literal["orchestration", "spatial_understanding", "image_generation", "image_editing"]
ImageProtocol = Literal["openai_images", "aigc_media"]
ImageParameterProfile = Literal["generic", "gpt_image_2", "nano_banana_pro", "seedream_5_pro"]
class ModelPoolItem(BaseModel):
@@ -76,6 +78,9 @@ class ModelPoolItem(BaseModel):
chat_path: str = "/chat/completions"
image_generation_path: str = "/images/generations"
image_edit_path: str = "/images/edits"
image_status_path: str = "/v1/media/status"
image_protocol: ImageProtocol = "openai_images"
image_parameter_profile: ImageParameterProfile = "generic"
enabled: bool = True
@@ -233,6 +238,37 @@ class EncryptedSettingsStore:
document["tests"].pop("ai_models", None)
return True
@staticmethod
def _migrate_model_protocol_fields(document: dict[str, Any]) -> bool:
changed = False
known_profiles = {
"gpt-image-2": "gpt_image_2",
"gemini-3-pro-image-preview": "nano_banana_pro",
"doubao-seedream-5-0-pro-260628": "seedream_5_pro",
}
for item in document.get("values", {}).get("model_pool", []):
if not isinstance(item, dict):
continue
defaults = {
"image_status_path": "/v1/media/status",
"image_protocol": "openai_images",
"image_parameter_profile": "generic",
}
host = urlparse(str(item.get("base_url", ""))).hostname or ""
profile = known_profiles.get(str(item.get("model_id", "")))
if host == "api.lk888.ai" and profile and "image_generation" in item.get("capabilities", []):
defaults.update(
image_protocol="aigc_media",
image_parameter_profile=profile,
image_generation_path="/v1/media/generate",
image_status_path="/v1/media/status",
)
for key, value in defaults.items():
if key not in item:
item[key] = value
changed = True
return changed
def load_document(self) -> dict[str, Any]:
if not self.data_path.exists():
return {"values": self.defaults(), "tests": {}}
@@ -243,7 +279,9 @@ class EncryptedSettingsStore:
raise RuntimeError("运行期配置无法解密,请检查主密钥是否发生变化。") from exc
document.setdefault("values", {})
document.setdefault("tests", {})
if self._migrate_legacy_models(document):
migrated = self._migrate_legacy_models(document)
migrated = self._migrate_model_protocol_fields(document) or migrated
if migrated:
self.save_document(document)
return document
@@ -718,14 +756,116 @@ async def test_model_pool_item(values: dict[str, Any], model_id: str) -> Runtime
message=f"模型“{name}”测试失败:账号模型列表中未找到 ID“{item['model_id']}”。",
details={"model_id": item["model_id"], "model_count": len(listed_ids)},
)
protocol_error = _validate_image_protocol(item)
if protocol_error:
return RuntimeSettingsTestResult(
target=target,
ok=False,
message=f"模型“{name}”的 ID 可用,但生图协议配置不正确:{protocol_error}",
details={"model_id": item["model_id"], "model_count": len(listed_ids)},
)
return RuntimeSettingsTestResult(
target=target,
ok=True,
message=f"模型“{name}”连接正常,模型 ID“{item['model_id']}”可用。",
message=f"模型“{name}”连接正常,模型 ID“{item['model_id']}”可用。此项不实际生图,不消耗额度。",
details={"model_id": item["model_id"], "model_count": len(listed_ids)},
)
def _validate_image_protocol(item: dict[str, Any]) -> str:
if "image_generation" not in item.get("capabilities", []):
return ""
host = urlparse(str(item.get("base_url", ""))).hostname or ""
known_aigc_models = {
"gpt-image-2": "gpt_image_2",
"gemini-3-pro-image-preview": "nano_banana_pro",
"doubao-seedream-5-0-pro-260628": "seedream_5_pro",
}
expected_profile = known_aigc_models.get(str(item.get("model_id", "")))
if host == "api.lk888.ai" and expected_profile:
if item.get("image_protocol", "openai_images") != "aigc_media":
return "聚合引擎上的该模型应选择“AIGC 异步媒体协议”"
generation_path = urlparse(_join_api_url(str(item["base_url"]), str(item.get("image_generation_path", "")))).path
status_path = urlparse(_join_api_url(str(item["base_url"]), str(item.get("image_status_path", "")))).path
if generation_path != "/v1/media/generate":
return "生图地址应解析为 /v1/media/generate"
if status_path != "/v1/media/status":
return "任务查询地址应解析为 /v1/media/status"
if item.get("image_parameter_profile", "generic") != expected_profile:
return "参数预设与模型不匹配,请重新选择对应的模型预设"
return ""
async def test_image_generation(values: dict[str, Any], model_id: str) -> RuntimeSettingsTestResult:
"""Run one real, billable image request to prove the configured protocol end to end."""
target = f"generation:{model_id}"
item = next(
(candidate for candidate in values.get("model_pool", []) if candidate.get("id") == model_id),
None,
)
if item is None:
return RuntimeSettingsTestResult(target=target, ok=False, message="模型不存在或已被删除。")
if "image_generation" not in item.get("capabilities", []):
return RuntimeSettingsTestResult(target=target, ok=False, message="该模型没有启用图像生成能力。")
protocol_error = _validate_image_protocol(item)
if protocol_error:
return RuntimeSettingsTestResult(target=target, ok=False, message=f"无法试生成:{protocol_error}")
profile = item.get("image_parameter_profile", "generic")
trial_options: dict[str, Any] = {}
if profile == "gpt_image_2":
trial_options = {"size": "1024x1024", "quality": "low"}
elif profile == "nano_banana_pro":
trial_options = {"aspectRatio": "1:1", "imageSize": "1K"}
elif profile == "seedream_5_pro":
trial_options = {"aspect_ratio": "1:1", "size": "1K"}
try:
from app.integrations.openai_compatible import OpenAICompatibleGateway
result = await OpenAICompatibleGateway(values).generate_image(
"极简室内材质测试图:一个米白色立方体放在浅灰背景中,无文字",
model_instance_id=model_id,
**trial_options,
)
output = _find_generated_image(result)
if not output:
return RuntimeSettingsTestResult(
target=target,
ok=False,
message="请求已结束,但响应中没有找到图片地址或图片数据。",
)
details = {"model_id": item["model_id"], "output_kind": output[0]}
if output[0] == "url":
details["result_url"] = output[1]
return RuntimeSettingsTestResult(
target=target,
ok=True,
message=f"模型“{item.get('name', item['model_id'])}”已完成一次真实生图,端到端配置可用。",
details=details,
)
except (httpx.HTTPError, RuntimeError, ValueError) as exc:
return RuntimeSettingsTestResult(target=target, ok=False, message=f"真实生图失败:{exc}")
def _find_generated_image(payload: dict[str, Any]) -> tuple[str, str] | None:
if payload.get("result_url"):
return "url", str(payload["result_url"])
for row in payload.get("data", []) if isinstance(payload.get("data"), list) else []:
if not isinstance(row, dict):
continue
if row.get("url"):
return "url", str(row["url"])
if row.get("b64_json"):
return "inline", ""
for candidate in payload.get("candidates", []) if isinstance(payload.get("candidates"), list) else []:
parts = candidate.get("content", {}).get("parts", []) if isinstance(candidate, dict) else []
for part in parts:
if isinstance(part, dict) and (part.get("inlineData", {}).get("data") or part.get("inline_data", {}).get("data")):
return "inline", ""
return None
async def _test_gpu(values: dict[str, Any]) -> RuntimeSettingsTestResult:
mode = values.get("gpu_mode", "disabled")
if mode == "disabled":
+112
View File
@@ -0,0 +1,112 @@
import pytest
from app.integrations.openai_compatible import OpenAICompatibleGateway
def image_values(profile: str, model_id: str) -> dict:
return {
"model_pool": [
{
"id": "image",
"name": "测试生图",
"model_id": model_id,
"category": "multimodal",
"provider": "lingke",
"base_url": "https://api.lk888.ai",
"api_key": "secret",
"capabilities": ["image_generation"],
"models_path": "/v1/models",
"image_protocol": "aigc_media",
"image_parameter_profile": profile,
"image_generation_path": "/v1/media/generate",
"image_status_path": "/v1/media/status",
"enabled": True,
}
],
"image_routing_mode": "manual",
"image_model_id": "image",
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("profile", "model_id", "options", "expected_params"),
[
("gpt_image_2", "gpt-image-2", {"size": "1024x1024", "quality": "low"}, {"size": "1024x1024", "quality": "low"}),
("nano_banana_pro", "gemini-3-pro-image-preview", {"aspect_ratio": "1:1", "size": "1K"}, {"aspectRatio": "1:1", "imageSize": "1K"}),
("seedream_5_pro", "doubao-seedream-5-0-pro-260628", {"aspectRatio": "1:1", "imageSize": "1K"}, {"aspect_ratio": "1:1", "size": "1K"}),
],
)
async def test_aigc_media_profiles_send_documented_parameter_names(
monkeypatch,
profile: str,
model_id: str,
options: dict,
expected_params: dict,
) -> None:
requests: list[dict] = []
class FakeResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return {"data": [{"url": "https://cdn.example.com/test.png"}]}
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *args) -> None:
return None
async def post(self, url, **kwargs) -> FakeResponse:
requests.append({"url": url, **kwargs})
return FakeResponse()
monkeypatch.setattr("app.integrations.openai_compatible.httpx.AsyncClient", lambda **kwargs: FakeClient())
result = await OpenAICompatibleGateway(image_values(profile, model_id)).generate_image("测试", **options)
assert result["data"][0]["url"].endswith("test.png")
assert requests[0]["url"] == "https://api.lk888.ai/v1/media/generate"
assert requests[0]["json"] == {"model": model_id, "prompt": "测试", "params": expected_params}
@pytest.mark.asyncio
async def test_aigc_media_polls_task_until_result_url(monkeypatch) -> None:
class FakeResponse:
def __init__(self, payload: dict) -> None:
self.payload = payload
def raise_for_status(self) -> None:
return None
def json(self) -> dict:
return self.payload
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *args) -> None:
return None
async def post(self, *args, **kwargs) -> FakeResponse:
return FakeResponse({"task_id": 123})
async def get(self, url, **kwargs) -> FakeResponse:
assert url == "https://api.lk888.ai/v1/media/status"
assert kwargs["params"] == {"task_id": 123}
return FakeResponse({"task_id": 123, "state": "success", "is_final": True, "result_url": "https://cdn.example.com/final.png"})
async def no_sleep(*args) -> None:
return None
monkeypatch.setattr("app.integrations.openai_compatible.httpx.AsyncClient", lambda **kwargs: FakeClient())
monkeypatch.setattr("app.integrations.openai_compatible.asyncio.sleep", no_sleep)
result = await OpenAICompatibleGateway(image_values("seedream_5_pro", "doubao-seedream-5-0-pro-260628")).generate_image("测试")
assert result["state"] == "success"
assert result["data"] == [{"url": "https://cdn.example.com/final.png"}]