78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
"""从中转站（OpenAI 兼容 /v1/models）取可用模型列表。
|
||
|
||
只返回模型 id，不返回、不记录任何密钥：key 只进请求头，
|
||
绝不进返回值、日志或缓存。列表缓存 1 小时（中转模型目录
|
||
很少变），中转挂了时调用方拿到缓存；无缓存又取不到时报
|
||
RelayError，由端点翻成 502。
|
||
|
||
base_url：环境变量 XCM_RELAY_BASE_URL，缺省
|
||
https://xcmapi.org/v1（与 codex-home 里 relay 配置的
|
||
base_url 一致）。key：环境变量 XCM_API_KEY（与
|
||
config.xcmapi.example.toml 的 env_key 一致）。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
|
||
DEFAULT_BASE_URL = "https://xcmapi.org/v1"
|
||
CACHE_TTL_SECONDS = 3600
|
||
|
||
|
||
class RelayError(RuntimeError):
|
||
pass
|
||
|
||
|
||
_cache = {"models": None, "updated_at": 0.0}
|
||
|
||
|
||
def _default_fetcher(url, headers, timeout):
|
||
req = urllib.request.Request(url, headers=headers, method="GET")
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
return json.loads(resp.read().decode("utf-8"))
|
||
except urllib.error.HTTPError as exc:
|
||
raise RelayError("relay models HTTP %s" % exc.code)
|
||
except (urllib.error.URLError, ValueError, TimeoutError) as exc:
|
||
raise RelayError("relay models fetch failed: %s" % exc)
|
||
|
||
|
||
def get_relay_models(fetcher=None, now=None):
|
||
"""返回 (models, cached, updated_at)：models 为 id 字符串列表。
|
||
|
||
fetcher 可注入（单测用）：签名 fetcher(url, headers, timeout)
|
||
-> 已解析的 JSON dict。缓存命中时不调 fetcher。
|
||
"""
|
||
now = time.time() if now is None else now
|
||
if (_cache["models"] is not None
|
||
and now - _cache["updated_at"] < CACHE_TTL_SECONDS):
|
||
return list(_cache["models"]), True, _cache["updated_at"]
|
||
base_url = os.environ.get("XCM_RELAY_BASE_URL") or DEFAULT_BASE_URL
|
||
url = base_url.rstrip("/") + "/models"
|
||
if fetcher is not None:
|
||
# 单测注入：不带 key、不碰环境，生产路径仍走下面强制 key。
|
||
data = fetcher(url, {}, 10)
|
||
else:
|
||
api_key = os.environ.get("XCM_API_KEY")
|
||
if not api_key:
|
||
raise RelayError("relay key not configured")
|
||
data = _default_fetcher(
|
||
url, {"Authorization": "Bearer " + api_key}, 10)
|
||
if not isinstance(data, dict) or not isinstance(data.get("data"), list):
|
||
raise RelayError("relay models unexpected response")
|
||
models = []
|
||
for item in data["data"]:
|
||
model_id = item.get("id") if isinstance(item, dict) else None
|
||
if isinstance(model_id, str) and model_id and model_id not in models:
|
||
models.append(model_id)
|
||
_cache["models"] = list(models)
|
||
_cache["updated_at"] = now
|
||
return list(models), False, now
|
||
|
||
|
||
def _reset_cache_for_tests():
|
||
_cache["models"] = None
|
||
_cache["updated_at"] = 0.0
|