Initial project import

2026-09-18 03:46:58 +08:00
commit b600da4b76
209 changed files with 45904 additions and 0 deletions
@@ -0,0 +1,182 @@
+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:  # noqa: BLE001 - try the next account for any upstream failure
+                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:  # noqa: BLE001 - retain cached models on upstream failure
+                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)