dsh-plugin-image-workbench 1.1.5 → 1.1.6

package/vendor/chatgpt2api/source/services/model_service.py

@@ -1,182 +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)
+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)