187 lines
8.0 KiB
Python
187 lines
8.0 KiB
Python
"""Authenticated model discovery shared by refresh and connection checks."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from dataclasses import dataclass
|
||
from datetime import datetime
|
||
from urllib.parse import urlsplit, urlunsplit
|
||
|
||
import httpx
|
||
|
||
CATALOG_TIMEOUT_SECONDS = 30
|
||
CATALOG_MAX_PAGES = 100
|
||
CATALOG_PAGE_SIZE = 100
|
||
CATALOG_REQUEST_ATTEMPTS = 2
|
||
CATALOG_RETRY_DELAY_SECONDS = 0.2
|
||
SUPPORTED_PROVIDER_APIS = {
|
||
"openai-completions",
|
||
"openai-responses",
|
||
"anthropic-messages",
|
||
"ollama-generate",
|
||
}
|
||
|
||
|
||
class LLMProviderCatalogError(RuntimeError):
|
||
"""A safe, user-facing catalog failure with no upstream response or credentials."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ModelCatalog:
|
||
url: str
|
||
models: list[str]
|
||
|
||
|
||
def model_catalog_url(provider: str, base_url: str, provider_api: str) -> str:
|
||
parts = urlsplit(base_url.strip())
|
||
if parts.scheme not in {"http", "https"} or not parts.hostname:
|
||
raise LLMProviderCatalogError("请填写有效的 HTTP(S) 模型基础地址。")
|
||
if parts.username or parts.password or parts.query or parts.fragment:
|
||
raise LLMProviderCatalogError("模型基础地址不能包含账号、密码、查询参数或片段。")
|
||
if provider_api not in SUPPORTED_PROVIDER_APIS:
|
||
raise LLMProviderCatalogError("当前接口协议不支持模型目录查询。")
|
||
path = parts.path.rstrip("/")
|
||
if provider_api == "ollama-generate":
|
||
path = path.removesuffix("/api").removesuffix("/v1") + "/api/tags"
|
||
elif provider == "alibaba" and parts.hostname.endswith(".aliyuncs.com"):
|
||
path = "/api/v1/models"
|
||
else:
|
||
if not path or (provider_api == "anthropic-messages" and path.endswith("/anthropic")):
|
||
path += "/v1"
|
||
path += "/models"
|
||
return urlunsplit((parts.scheme, parts.netloc, path, "", ""))
|
||
|
||
|
||
def catalog_error_message(exc: Exception) -> str:
|
||
if isinstance(exc, LLMProviderCatalogError):
|
||
return str(exc)
|
||
if isinstance(exc, (httpx.TimeoutException, TimeoutError)):
|
||
return "模型目录请求超时，请检查网络后重试。"
|
||
if isinstance(exc, httpx.HTTPStatusError):
|
||
code = exc.response.status_code
|
||
messages = {
|
||
401: "API Key 验证失败，请检查当前供应商的凭证。",
|
||
403: "当前 API Key 无权访问该模型目录，请检查账号权限和服务地域。",
|
||
404: "模型目录接口不存在，请检查基础地址、地域和接口协议。",
|
||
429: "供应商请求限流，请稍后重试。",
|
||
}
|
||
return messages.get(code, f"供应商模型目录返回 HTTP {code}，请稍后重试。")
|
||
if isinstance(exc, httpx.RequestError):
|
||
return "无法连接模型目录，请检查基础地址和网络。"
|
||
return "模型目录响应无效，请稍后重试。"
|
||
|
||
|
||
def _model_rows(payload: object) -> tuple[list[dict[str, object]], dict[str, object]]:
|
||
if not isinstance(payload, dict):
|
||
raise LLMProviderCatalogError("供应商返回了无效的模型目录。")
|
||
envelope = payload.get("output", payload)
|
||
if not isinstance(envelope, dict):
|
||
raise LLMProviderCatalogError("供应商返回了无效的模型目录。")
|
||
rows = envelope.get("data", envelope.get("models"))
|
||
if not isinstance(rows, list):
|
||
raise LLMProviderCatalogError("供应商响应中没有模型列表。")
|
||
if any(not isinstance(row, dict) for row in rows):
|
||
raise LLMProviderCatalogError("供应商返回了无效的模型条目。")
|
||
return rows, envelope
|
||
|
||
|
||
def _model_id(row: dict[str, object]) -> str:
|
||
value = row.get("id") or row.get("model") or row.get("name")
|
||
if not isinstance(value, str) or not value.strip():
|
||
raise LLMProviderCatalogError("供应商返回了缺少 ID 的模型条目。")
|
||
return value.strip()
|
||
|
||
|
||
def _model_date(row: dict[str, object]) -> float:
|
||
value = row.get("created_at") or row.get("published_time") or row.get("created")
|
||
if isinstance(value, (int, float)):
|
||
return float(value)
|
||
if isinstance(value, str):
|
||
try:
|
||
return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
|
||
except ValueError:
|
||
pass
|
||
return 0
|
||
|
||
|
||
async def _get_catalog_page(
|
||
client: httpx.AsyncClient,
|
||
url: str,
|
||
headers: dict[str, str],
|
||
params: dict[str, str | int],
|
||
) -> object:
|
||
for attempt in range(CATALOG_REQUEST_ATTEMPTS):
|
||
try:
|
||
response = await client.get(url, headers=headers, params=params)
|
||
response.raise_for_status()
|
||
return response.json()
|
||
except httpx.HTTPStatusError as exc:
|
||
if attempt or exc.response.status_code not in {502, 503, 504}:
|
||
raise
|
||
except httpx.TransportError:
|
||
if attempt:
|
||
raise
|
||
await asyncio.sleep(CATALOG_RETRY_DELAY_SECONDS)
|
||
raise LLMProviderCatalogError("模型目录请求失败。")
|
||
|
||
|
||
async def fetch_model_catalog(
|
||
provider: str,
|
||
base_url: str,
|
||
provider_api: str,
|
||
api_key: str = "",
|
||
anthropic_version: str = "2023-06-01",
|
||
timeout_seconds: int = CATALOG_TIMEOUT_SECONDS,
|
||
) -> ModelCatalog:
|
||
url = model_catalog_url(provider, base_url, provider_api)
|
||
public_catalog = provider in {"opencode-go", "openrouter"}
|
||
if not api_key and provider_api != "ollama-generate" and not public_catalog:
|
||
raise LLMProviderCatalogError("请先配置当前供应商的 API Key，再刷新模型列表。")
|
||
headers = {"User-Agent": "Planet/1.0", "Accept": "application/json"}
|
||
if provider_api == "anthropic-messages" and provider not in {
|
||
"opencode-go",
|
||
"openrouter",
|
||
"alibaba",
|
||
"moonshotai",
|
||
}:
|
||
headers.update({"x-api-key": api_key, "anthropic-version": anthropic_version})
|
||
elif api_key:
|
||
headers["Authorization"] = f"Bearer {api_key}"
|
||
native_dashscope = provider == "alibaba" and urlsplit(url).hostname.endswith(".aliyuncs.com")
|
||
params: dict[str, str | int] = {}
|
||
if native_dashscope:
|
||
params = {"page_no": 1, "page_size": CATALOG_PAGE_SIZE, "capabilities": "TG"}
|
||
rows_by_id: dict[str, dict[str, object]] = {}
|
||
timeout = max(1, min(timeout_seconds, CATALOG_TIMEOUT_SECONDS))
|
||
# Bound the complete pagination/retry cycle, not just each individual request.
|
||
async with asyncio.timeout(timeout), httpx.AsyncClient(timeout=timeout) as client:
|
||
for page in range(CATALOG_MAX_PAGES):
|
||
payload = await _get_catalog_page(client, url, headers, params)
|
||
rows, envelope = _model_rows(payload)
|
||
previous_count = len(rows_by_id)
|
||
for row in rows:
|
||
rows_by_id[_model_id(row)] = row
|
||
has_more = envelope.get("has_more") is True
|
||
if native_dashscope:
|
||
total = envelope.get("total")
|
||
if not isinstance(total, int) or total < 0:
|
||
raise LLMProviderCatalogError("供应商返回了无效的模型目录分页信息。")
|
||
has_more = len(rows_by_id) < total
|
||
params["page_no"] = page + 2
|
||
elif has_more:
|
||
cursor = envelope.get("last_id")
|
||
if not isinstance(cursor, str) or not cursor or cursor == params.get("after_id"):
|
||
raise LLMProviderCatalogError("供应商返回了无效的模型目录分页信息。")
|
||
params["after_id"] = cursor
|
||
if not has_more:
|
||
models = sorted(
|
||
rows_by_id, key=lambda key: _model_date(rows_by_id[key]), reverse=True
|
||
)
|
||
# An empty Ollama catalog is valid: no models have been installed yet.
|
||
if not models and provider_api != "ollama-generate":
|
||
raise LLMProviderCatalogError("供应商返回了空模型目录，已保留上次模型列表。")
|
||
return ModelCatalog(url=url, models=models)
|
||
if len(rows_by_id) == previous_count:
|
||
raise LLMProviderCatalogError("供应商模型目录分页没有进展，已保留上次模型列表。")
|
||
raise LLMProviderCatalogError("供应商模型目录分页超过限制，已保留上次模型列表。")
|