196 lines
7.8 KiB
Python
196 lines
7.8 KiB
Python
import logging
|
|
import httpx
|
|
|
|
from .channel_config import load_channel_groups, resolve_channel_config, save_channel_groups
|
|
|
|
logger = logging.getLogger("aichat_web_api")
|
|
|
|
try:
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
|
|
_HAS_SERVER = PromptServer.instance is not None
|
|
except Exception as e:
|
|
logger.warning(f"[aichat] 未能导入 PromptServer / aiohttp，接口已禁用: {e}")
|
|
_HAS_SERVER = False
|
|
|
|
BASE = "/aichat/api/"
|
|
|
|
|
|
def _build_proxies(proxy_http: str, proxy_https: str):
|
|
proxies = {}
|
|
if proxy_http and proxy_http.strip():
|
|
proxies["http://"] = proxy_http.strip()
|
|
if proxy_https and proxy_https.strip():
|
|
proxies["https://"] = proxy_https.strip()
|
|
return proxies or None
|
|
|
|
|
|
def _make_async_client(timeout: float, proxies, **kwargs) -> "httpx.AsyncClient":
|
|
if proxies:
|
|
return httpx.AsyncClient(proxies=proxies, timeout=timeout, **kwargs)
|
|
return httpx.AsyncClient(timeout=timeout, **kwargs)
|
|
|
|
|
|
async def _read_json(request):
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
data = {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
async def _read_common_params(request):
|
|
data = await _read_json(request)
|
|
try:
|
|
timeout = float(data.get("timeout") or 60.0)
|
|
except (TypeError, ValueError):
|
|
timeout = 60.0
|
|
return {
|
|
"provider": (data.get("provider") or "").strip().lower(),
|
|
"channel_group": (data.get("channel_group") or "").strip(),
|
|
"api_key": (data.get("api_key") or "").strip(),
|
|
"base_url": (data.get("base_url") or "").strip(),
|
|
"proxy_http": data.get("proxy_http") or "",
|
|
"proxy_https": data.get("proxy_https") or "",
|
|
"timeout": timeout,
|
|
}
|
|
|
|
|
|
def _extract_openai_models(body) -> list:
|
|
items = []
|
|
if isinstance(body, dict):
|
|
items = body.get("data") or body.get("models") or []
|
|
elif isinstance(body, list):
|
|
items = body
|
|
|
|
out = []
|
|
for it in items:
|
|
if isinstance(it, str):
|
|
out.append(it)
|
|
elif isinstance(it, dict):
|
|
mid = it.get("id") or it.get("name") or it.get("model")
|
|
if mid:
|
|
out.append(str(mid))
|
|
return sorted(set(out))
|
|
|
|
|
|
if _HAS_SERVER:
|
|
|
|
@PromptServer.instance.routes.post(BASE + "groups/get")
|
|
async def _aichat_groups_get(request):
|
|
data = await _read_json(request)
|
|
provider = (data.get("provider") or "").strip().lower()
|
|
try:
|
|
groups = load_channel_groups(provider)
|
|
return web.json_response({"groups": groups, "error": None})
|
|
except Exception as e:
|
|
logger.error(f"[aichat] groups/get 失败: {type(e).__name__} - {e}")
|
|
return web.json_response({"groups": [], "error": f"{type(e).__name__}: {e}"})
|
|
|
|
@PromptServer.instance.routes.post(BASE + "groups/save")
|
|
async def _aichat_groups_save(request):
|
|
data = await _read_json(request)
|
|
provider = (data.get("provider") or "").strip().lower()
|
|
groups = data.get("groups") or []
|
|
try:
|
|
saved = save_channel_groups(provider, groups)
|
|
return web.json_response({"groups": saved, "error": None})
|
|
except Exception as e:
|
|
logger.error(f"[aichat] groups/save 失败: {type(e).__name__} - {e}")
|
|
return web.json_response({"groups": [], "error": f"{type(e).__name__}: {e}"})
|
|
|
|
@PromptServer.instance.routes.post(BASE + "openai/get_models")
|
|
async def _aichat_openai_get_models(request):
|
|
p = await _read_common_params(request)
|
|
try:
|
|
api_key, base_url, _ = resolve_channel_config(
|
|
"openai",
|
|
channel_group=p["channel_group"],
|
|
api_key=p["api_key"],
|
|
base_url=p["base_url"],
|
|
)
|
|
except Exception as e:
|
|
return web.json_response({"models": [], "error": str(e)})
|
|
|
|
if not api_key:
|
|
return web.json_response({"models": [], "error": "API Key 不能为空"})
|
|
if not base_url:
|
|
return web.json_response({"models": [], "error": "Base URL 不能为空"})
|
|
|
|
if not base_url.endswith("/"):
|
|
base_url += "/"
|
|
proxies = _build_proxies(p["proxy_http"], p["proxy_https"])
|
|
url = base_url + "models"
|
|
headers = {"Authorization": f"Bearer {api_key}"}
|
|
try:
|
|
async with _make_async_client(p["timeout"], proxies) as client:
|
|
resp = await client.get(url, headers=headers)
|
|
if resp.status_code != 200:
|
|
snippet = resp.text[:300] if resp.text else ""
|
|
return web.json_response({"models": [], "error": f"HTTP {resp.status_code}: {snippet}"})
|
|
models = _extract_openai_models(resp.json())
|
|
logger.info(f"[aichat] OpenAI 获取到 {len(models)} 个模型 ({base_url})")
|
|
return web.json_response({"models": models, "error": None})
|
|
except Exception as e:
|
|
logger.error(f"[aichat] OpenAI get_models 失败: {type(e).__name__} - {e}")
|
|
return web.json_response({"models": [], "error": f"{type(e).__name__}: {e}"})
|
|
|
|
@PromptServer.instance.routes.post(BASE + "gemini/get_models")
|
|
async def _aichat_gemini_get_models(request):
|
|
p = await _read_common_params(request)
|
|
try:
|
|
api_key, base_url, _ = resolve_channel_config(
|
|
"gemini",
|
|
channel_group=p["channel_group"],
|
|
api_key=p["api_key"],
|
|
base_url=p["base_url"],
|
|
)
|
|
except Exception as e:
|
|
return web.json_response({"models": [], "error": str(e)})
|
|
|
|
if not api_key:
|
|
return web.json_response({"models": [], "error": "API Key 不能为空"})
|
|
if not base_url:
|
|
return web.json_response({"models": [], "error": "Base URL 不能为空"})
|
|
|
|
base_url = base_url.rstrip("/")
|
|
proxies = _build_proxies(p["proxy_http"], p["proxy_https"])
|
|
try:
|
|
async with _make_async_client(
|
|
p["timeout"], proxies, base_url=base_url, params={"key": api_key}
|
|
) as client:
|
|
models = []
|
|
page_token = None
|
|
for _ in range(20):
|
|
params = {"pageSize": 1000}
|
|
if page_token:
|
|
params["pageToken"] = page_token
|
|
resp = await client.get("/v1beta/models", params=params)
|
|
if resp.status_code != 200:
|
|
snippet = resp.text[:300] if resp.text else ""
|
|
return web.json_response({"models": [], "error": f"HTTP {resp.status_code}: {snippet}"})
|
|
body = resp.json()
|
|
for m in body.get("models", []):
|
|
name = m.get("name", "")
|
|
methods = m.get("supportedGenerationMethods") or []
|
|
if methods and not any(x in methods for x in ("generateContent", "streamGenerateContent")):
|
|
continue
|
|
if name.startswith("models/"):
|
|
name = name[len("models/"):]
|
|
if name:
|
|
models.append(name)
|
|
page_token = body.get("nextPageToken")
|
|
if not page_token:
|
|
break
|
|
models = sorted(set(models))
|
|
logger.info(f"[aichat] Gemini 获取到 {len(models)} 个模型 ({base_url})")
|
|
return web.json_response({"models": models, "error": None})
|
|
except Exception as e:
|
|
logger.error(f"[aichat] Gemini get_models 失败: {type(e).__name__} - {e}")
|
|
return web.json_response({"models": [], "error": f"{type(e).__name__}: {e}"})
|
|
|
|
print(
|
|
f"[aichat] 接口已注册: POST {BASE}groups/get | POST {BASE}groups/save | POST {BASE}openai/get_models | POST {BASE}gemini/get_models"
|
|
)
|