| from __future__ import annotations |
|
|
| import time |
| from collections.abc import Callable |
| from concurrent.futures import ThreadPoolExecutor |
| from dataclasses import dataclass |
| from threading import RLock |
| from typing import Any |
|
|
| from services.account_service import AccountService, account_service |
| from services.openai_backend_api import OpenAIBackendAPI |
| from utils.log import logger |
|
|
|
|
| @dataclass(frozen=True) |
| class ModelRoute: |
| account_types: frozenset[str] |
| allow_anonymous: bool = False |
|
|
|
|
| class ModelUnavailableError(RuntimeError): |
| pass |
|
|
|
|
| class ModelCatalogService: |
| """Caches the model catalogs advertised to each active account type.""" |
|
|
| def __init__( |
| self, |
| accounts: AccountService, |
| *, |
| backend_factory: Callable[..., Any] = OpenAIBackendAPI, |
| cache_ttl_seconds: float = 300, |
| clock: Callable[[], float] = time.monotonic, |
| ) -> None: |
| self._accounts = accounts |
| self._backend_factory = backend_factory |
| self._cache_ttl_seconds = max(1.0, float(cache_ttl_seconds)) |
| self._clock = clock |
| self._lock = RLock() |
| self._expires_at = 0.0 |
| self._account_signature: tuple[tuple[str, int], ...] = () |
| self._anonymous_models: dict[str, dict[str, Any]] = {} |
| self._models_by_account_type: dict[str, dict[str, dict[str, Any]]] = {} |
|
|
| @staticmethod |
| def _model_map(result: object) -> dict[str, dict[str, Any]]: |
| if not isinstance(result, dict) or not isinstance(result.get("data"), list): |
| raise TypeError("upstream model response has no data list") |
| models: dict[str, dict[str, Any]] = {} |
| for item in result["data"]: |
| if not isinstance(item, dict): |
| continue |
| model_id = str(item.get("id") or "").strip() |
| if model_id and model_id not in models: |
| models[model_id] = dict(item) |
| return models |
|
|
| def _active_accounts_by_type(self) -> dict[str, list[str]]: |
| groups: dict[str, list[str]] = {} |
| for account in self._accounts.list_accounts(): |
| if not isinstance(account, dict) or account.get("status") in {"禁用", "异常"}: |
| continue |
| access_token = str(account.get("access_token") or "").strip() |
| account_type = self._accounts._normalize_account_type(account.get("type")) |
| if access_token and account_type: |
| groups.setdefault(account_type, []).append(access_token) |
| return groups |
|
|
| @staticmethod |
| def _signature(groups: dict[str, list[str]]) -> tuple[tuple[str, int], ...]: |
| return tuple( |
| (account_type, len(tokens)) |
| for account_type, tokens in sorted(groups.items()) |
| ) |
|
|
| def _fetch_models(self, access_token: str = "") -> dict[str, dict[str, Any]]: |
| backend = self._backend_factory(access_token=access_token) |
| try: |
| return self._model_map(backend.list_models()) |
| finally: |
| backend.close() |
|
|
| def _fetch_account_type_models( |
| self, |
| account_type: str, |
| access_tokens: list[str], |
| ) -> dict[str, dict[str, Any]] | None: |
| attempted_tokens: set[str] = set() |
| last_error: Exception | None = None |
| for access_token in access_tokens: |
| try: |
| resolved_token = self._accounts.refresh_access_token( |
| access_token, |
| event="list_models", |
| ) or access_token |
| if resolved_token in attempted_tokens: |
| continue |
| attempted_tokens.add(resolved_token) |
| return self._fetch_models(resolved_token) |
| except Exception as exc: |
| last_error = exc |
| if last_error is not None: |
| logger.warning({ |
| "event": "model_catalog_account_type_failed", |
| "account_type": account_type, |
| "error_type": type(last_error).__name__, |
| }) |
| return None |
|
|
| def _refresh(self, groups: dict[str, list[str]], signature: tuple[tuple[str, int], ...]) -> None: |
| models_by_account_type: dict[str, dict[str, dict[str, Any]]] = {} |
| with ThreadPoolExecutor(max_workers=min(4, len(groups) + 1)) as executor: |
| anonymous_future = executor.submit(self._fetch_models) |
| account_futures = { |
| account_type: executor.submit( |
| self._fetch_account_type_models, |
| account_type, |
| access_tokens, |
| ) |
| for account_type, access_tokens in groups.items() |
| } |
| try: |
| anonymous_models = anonymous_future.result() |
| except Exception as exc: |
| logger.warning({ |
| "event": "model_catalog_anonymous_failed", |
| "error_type": type(exc).__name__, |
| }) |
| anonymous_models = self._anonymous_models |
|
|
| for account_type, future in account_futures.items(): |
| models = future.result() |
| if models is not None: |
| models_by_account_type[account_type] = models |
| elif account_type in self._models_by_account_type: |
| models_by_account_type[account_type] = self._models_by_account_type[account_type] |
|
|
| self._anonymous_models = anonymous_models |
| self._models_by_account_type = models_by_account_type |
| self._account_signature = signature |
| self._expires_at = self._clock() + self._cache_ttl_seconds |
|
|
| def _ensure_catalog(self) -> None: |
| groups = self._active_accounts_by_type() |
| signature = self._signature(groups) |
| with self._lock: |
| if signature == self._account_signature and self._clock() < self._expires_at: |
| return |
| self._refresh(groups, signature) |
|
|
| def list_models(self) -> dict[str, Any]: |
| self._ensure_catalog() |
| with self._lock: |
| union: dict[str, dict[str, Any]] = { |
| model_id: dict(item) |
| for model_id, item in self._anonymous_models.items() |
| } |
| for account_type in sorted(self._models_by_account_type): |
| for model_id, item in self._models_by_account_type[account_type].items(): |
| union.setdefault(model_id, dict(item)) |
| return { |
| "object": "list", |
| "data": [union[model_id] for model_id in sorted(union)], |
| } |
|
|
| def route_for_model(self, model: str) -> ModelRoute: |
| model = str(model or "").strip() |
| self._ensure_catalog() |
| with self._lock: |
| account_types = frozenset( |
| account_type |
| for account_type, models in self._models_by_account_type.items() |
| if model in models |
| ) |
| return ModelRoute( |
| account_types=account_types, |
| allow_anonymous=model in self._anonymous_models, |
| ) |
|
|
|
|
| model_catalog_service = ModelCatalogService(account_service) |
|
|
