feat: add configurable AI model pool

This commit is contained in:
Codex
2026-08-01 23:11:44 +08:00
parent 020b9596dc
commit 4430149ee4
13 changed files with 1660 additions and 273 deletions
@@ -4,48 +4,70 @@ from urllib.parse import urljoin
import httpx
from app.integrations.model_router import ModelRouter
class OpenAICompatibleGateway:
"""One adapter for OpenAI, lk666.ai and other OpenAI-compatible gateways."""
"""Adapter for independently configured OpenAI-compatible model-pool items."""
def __init__(self, values: Mapping[str, Any]) -> None:
self.values = values
self.router = ModelRouter(values)
def _url(self, path_key: str, fallback: str) -> str:
base_url = str(self.values.get("ai_base_url", "")).rstrip("/") + "/"
path = str(self.values.get(path_key, fallback)).lstrip("/")
@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)
@property
def headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self.values.get('ai_api_key', '')}",
"Content-Type": "application/json",
}
@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) -> list[str]:
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("ai_models_path", "/models"), headers=self.headers)
response = await client.get(
self._url(item, "models_path", "/models"),
headers=self._headers(item),
)
response.raise_for_status()
payload = response.json()
return [item["id"] for item in payload.get("data", []) if isinstance(item, dict) and item.get("id")]
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) -> dict[str, Any]:
model_key = "vision_model" if vision else "orchestrator_model"
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("ai_chat_path", "/chat/completions"),
headers=self.headers,
json={"model": self.values[model_key], "messages": messages},
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, **options: Any) -> dict[str, Any]:
payload = {"model": self.values["image_model"], "prompt": prompt, **options}
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)
payload = {"model": item["model_id"], "prompt": prompt, **options}
async with httpx.AsyncClient(timeout=180) as client:
response = await client.post(
self._url("ai_image_generation_path", "/images/generations"),
headers=self.headers,
self._url(item, "image_generation_path", "/images/generations"),
headers=self._headers(item),
json=payload,
)
response.raise_for_status()
@@ -58,19 +80,18 @@ class OpenAICompatibleGateway:
*,
filename: str = "image.png",
mask: bytes | None = None,
model_instance_id: str | None = None,
**options: Any,
) -> dict[str, Any]:
files: dict[str, tuple[str, bytes, str]] = {
"image": (filename, image, "image/png"),
}
item = self.router.choose("image", model_instance_id)
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": self.values["image_model"], "prompt": prompt, **options}
headers = {"Authorization": self.headers["Authorization"]}
data = {"model": item["model_id"], "prompt": prompt, **options}
async with httpx.AsyncClient(timeout=180) as client:
response = await client.post(
self._url("ai_image_edit_path", "/images/edits"),
headers=headers,
self._url(item, "image_edit_path", "/images/edits"),
headers=self._headers(item, json_content=False),
data=data,
files=files,
)