feat: upgrade image routing and reference workflows

Author: root

.gitignore

@@ -9,7 +9,10 @@ wheels/
 # Virtual environments
 .venv
 .idea
+.env
+.env.*
+token/
 data
 config.json
 
-docker-compose-local.yml
+docker-compose-local.yml

README.md

@@ -9,6 +9,7 @@ ChatGPT 图片生成代理与账号池管理面板，提供账号维护、额度
 - 批量导入和管理 `access_token`
 - 自动刷新账号邮箱、类型、图片额度、恢复时间
 - 轮询可用账号进行图片生成
+- `gpt-image-2` 账号优先级路由、自动降级、后台自动补池
 - 失效 Token 自动剔除
 - 提供 Web 后台管理账号和生成图片 
 
@@ -54,6 +55,7 @@ git clone git@github.com:basketikun/chatgpt2api.git
 cd chatgpt2api
 export CHATGPT2API_AUTH_KEY='replace-with-your-auth-key'
 export CHATGPT2API_PUBLIC_BASE_URL='http://37.114.42.229:9090'
+export CHATGPT2API_IMG2_PROBE_INTERVAL_SECONDS='900'
 docker compose up --build -d
 ```
 

config.example.json

@@ -3,5 +3,8 @@
   "public-base-url": "http://127.0.0.1:9090",
   "image-ttl-hours": 360,
   "image-cleanup-interval-seconds": 300,
+  "img2-probe-interval-seconds": 900,
+  "img2-probe-runs": 3,
+  "img2-probe-concurrency": 4,
   "tls-verify": true
 }

docker-compose.yml

@@ -9,6 +9,9 @@ services:
       CHATGPT2API_PUBLIC_BASE_URL: ${CHATGPT2API_PUBLIC_BASE_URL:-}
       CHATGPT2API_IMAGE_TTL_HOURS: ${CHATGPT2API_IMAGE_TTL_HOURS:-360}
       CHATGPT2API_IMAGE_CLEANUP_INTERVAL_SECONDS: ${CHATGPT2API_IMAGE_CLEANUP_INTERVAL_SECONDS:-300}
+      CHATGPT2API_IMG2_PROBE_INTERVAL_SECONDS: ${CHATGPT2API_IMG2_PROBE_INTERVAL_SECONDS:-900}
+      CHATGPT2API_IMG2_PROBE_RUNS: ${CHATGPT2API_IMG2_PROBE_RUNS:-3}
+      CHATGPT2API_IMG2_PROBE_CONCURRENCY: ${CHATGPT2API_IMG2_PROBE_CONCURRENCY:-4}
     ports:
       - "9090:80"
     volumes:

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.13"
 dependencies = [
     "curl-cffi>=0.15.0",
     "fastapi>=0.136.0",
+    "python-multipart>=0.0.20",
     "pybase64>=1.4.3",
     "uvicorn>=0.44.0",
 ]

services/account_service.py

@@ -30,6 +30,7 @@ def __init__(self, store_file: Path):
         self._lock = Lock()
         self._index = 0
         self._accounts = self._load_accounts()
+        self._backfill_missing_preferred_models()
 
     @staticmethod
     def _clean_token(value: Any) -> str:
@@ -59,6 +60,16 @@ def _is_image_account_available(account: dict) -> bool:
             return False
         return int(account.get("quota") or 0) > 0
 
+    @staticmethod
+    def _should_default_to_img1(account: dict) -> bool:
+        if not isinstance(account, dict):
+            return False
+        if account.get("preferred_models"):
+            return False
+        if str(account.get("status") or "").strip() in {"禁用", "异常"}:
+            return False
+        return int(account.get("quota") or 0) > 0
+
     def _decode_access_token_payload(self, access_token: str) -> dict[str, Any]:
         parts = self._clean_token(access_token).split(".")
         if len(parts) < 2:
@@ -138,6 +149,27 @@ def _normalize_account(self, item: dict) -> dict | None:
             ]
         else:
             normalized["preferred_models"] = []
+        if self._should_default_to_img1(normalized):
+            normalized["preferred_models"] = ["gpt-image-1"]
+        routing_priority = normalized.get("routing_priority")
+        if routing_priority is None or str(routing_priority).strip() == "":
+            routing_priority = 100 if "gpt-image-2" in normalized["preferred_models"] else 0
+        normalized["routing_priority"] = int(routing_priority)
+        recent_img2_results = normalized.get("recent_img2_results")
+        if isinstance(recent_img2_results, list):
+            normalized["recent_img2_results"] = [
+                self._clean_token(item)
+                for item in recent_img2_results
+                if self._clean_token(item) in {"gptimage1", "gptimage2"}
+            ]
+        else:
+            normalized["recent_img2_results"] = []
+        normalized["consecutive_img1_results"] = max(0, int(normalized.get("consecutive_img1_results") or 0))
+        normalized["auto_managed"] = bool(normalized.get("auto_managed") or False)
+        normalized["last_probe_at"] = self._clean_token(normalized.get("last_probe_at")) or None
+        normalized["last_probe_result"] = self._clean_token(normalized.get("last_probe_result")) or None
+        normalized["promotion_count"] = int(normalized.get("promotion_count") or 0)
+        normalized["demotion_count"] = int(normalized.get("demotion_count") or 0)
         normalized["restore_at"] = self._clean_token(normalized.get("restore_at")) or None
         normalized["success"] = int(normalized.get("success") or 0)
         normalized["fail"] = int(normalized.get("fail") or 0)
@@ -167,6 +199,20 @@ def _load_accounts(self) -> list[dict]:
             return []
         return [normalized for item in data if (normalized := self._normalize_account(item)) is not None]
 
+    def _backfill_missing_preferred_models(self) -> None:
+        changed = False
+        next_accounts: list[dict] = []
+        for item in self._accounts:
+            next_item = dict(item)
+            if self._should_default_to_img1(next_item):
+                next_item["preferred_models"] = ["gpt-image-1"]
+                next_item["routing_priority"] = int(next_item.get("routing_priority") or 0)
+                changed = True
+            next_accounts.append(next_item)
+        if changed:
+            self._accounts = next_accounts
+            self._save_accounts()
+
     def _save_accounts(self) -> None:
         self.store_file.parent.mkdir(parents=True, exist_ok=True)
         self.store_file.write_text(
@@ -218,6 +264,14 @@ def _public_items(self, accounts: list[dict]) -> list[dict]:
                 "limits_progress": account.get("limits_progress") or [],
                 "default_model_slug": account.get("default_model_slug"),
                 "preferred_models": account.get("preferred_models") or [],
+                "routing_priority": int(account.get("routing_priority") or 0),
+                "recent_img2_results": account.get("recent_img2_results") or [],
+                "consecutive_img1_results": int(account.get("consecutive_img1_results") or 0),
+                "auto_managed": bool(account.get("auto_managed") or False),
+                "last_probe_at": account.get("last_probe_at"),
+                "last_probe_result": account.get("last_probe_result"),
+                "promotion_count": int(account.get("promotion_count") or 0),
+                "demotion_count": int(account.get("demotion_count") or 0),
                 "restoreAt": account.get("restore_at"),
                 "success": int(account.get("success") or 0),
                 "fail": int(account.get("fail") or 0),
@@ -246,13 +300,19 @@ def next_token(self, excluded_tokens: set[str] | None = None, model: str | None
                 raise RuntimeError(f"No available tokens found in {self.store_file}")
             if model_name:
                 preferred = [
-                    token
+                    (item, token)
                     for item, token in candidates
                     if model_name in (item.get("preferred_models") or [])
                 ]
-                tokens = preferred or [token for _, token in candidates]
+                selected_candidates = preferred or candidates
             else:
-                tokens = [token for _, token in candidates]
+                selected_candidates = candidates
+            max_priority = max(int(item.get("routing_priority") or 0) for item, _ in selected_candidates)
+            tokens = [
+                token
+                for item, token in selected_candidates
+                if int(item.get("routing_priority") or 0) == max_priority
+            ]
             access_token = tokens[self._index % len(tokens)]
             self._index += 1
             return access_token
@@ -271,6 +331,20 @@ def list_accounts(self) -> list[dict]:
         with self._lock:
             return self._public_items(self._accounts)
 
+    def list_available_accounts(self, model: str | None = None) -> list[dict]:
+        model_name = self._clean_token(model)
+        with self._lock:
+            items = [
+                dict(item)
+                for item in self._accounts
+                if self._is_image_account_available(item)
+                and (not model_name or model_name in (item.get("preferred_models") or []))
+            ]
+        return items
+
+    def count_available_accounts(self, model: str | None = None) -> int:
+        return len(self.list_available_accounts(model=model))
+
     def list_limited_tokens(self) -> list[str]:
         with self._lock:
             return [
@@ -374,6 +448,119 @@ def mark_image_result(self, access_token: str, success: bool) -> dict | None:
             return dict(account)
         return None
 
+    def demote_to_img1(self, access_token: str, reason: str = "auto") -> dict | None:
+        access_token = self._clean_token(access_token)
+        if not access_token:
+            return None
+        with self._lock:
+            index = self._find_account_index(access_token)
+            if index < 0:
+                return None
+            next_item = dict(self._accounts[index])
+            next_item["preferred_models"] = ["gpt-image-1"]
+            next_item["routing_priority"] = 0
+            next_item["auto_managed"] = True
+            next_item["last_probe_result"] = self._clean_token(reason) or "auto_demoted"
+            next_item["demotion_count"] = int(next_item.get("demotion_count") or 0) + 1
+            account = self._normalize_account(next_item)
+            if account is None:
+                return None
+            self._accounts[index] = account
+            self._save_accounts()
+            return dict(account)
+
+    def promote_to_img2(self, access_token: str, priority: int = 100, reason: str = "probe_promoted") -> dict | None:
+        access_token = self._clean_token(access_token)
+        if not access_token:
+            return None
+        with self._lock:
+            index = self._find_account_index(access_token)
+            if index < 0:
+                return None
+            next_item = dict(self._accounts[index])
+            next_item["preferred_models"] = ["gpt-image-2"]
+            next_item["routing_priority"] = int(priority)
+            next_item["auto_managed"] = True
+            next_item["last_probe_at"] = datetime.now().isoformat(timespec="seconds")
+            next_item["last_probe_result"] = self._clean_token(reason) or "probe_promoted"
+            next_item["promotion_count"] = int(next_item.get("promotion_count") or 0) + 1
+            next_item["consecutive_img1_results"] = 0
+            next_item["recent_img2_results"] = []
+            account = self._normalize_account(next_item)
+            if account is None:
+                return None
+            self._accounts[index] = account
+            self._save_accounts()
+            return dict(account)
+
+    def update_probe_result(self, access_token: str, result: str) -> dict | None:
+        access_token = self._clean_token(access_token)
+        if not access_token:
+            return None
+        with self._lock:
+            index = self._find_account_index(access_token)
+            if index < 0:
+                return None
+            next_item = dict(self._accounts[index])
+            next_item["last_probe_at"] = datetime.now().isoformat(timespec="seconds")
+            next_item["last_probe_result"] = self._clean_token(result) or None
+            account = self._normalize_account(next_item)
+            if account is None:
+                return None
+            self._accounts[index] = account
+            self._save_accounts()
+            return dict(account)
+
+    def record_img2_outcome(
+            self,
+            access_token: str,
+            classification: str,
+            *,
+            recent_window: int = 6,
+            recent_img1_threshold: int = 4,
+            consecutive_img1_threshold: int = 3,
+    ) -> dict | None:
+        access_token = self._clean_token(access_token)
+        classification = self._clean_token(classification)
+        if not access_token or classification not in {"gptimage1", "gptimage2"}:
+            return None
+        with self._lock:
+            index = self._find_account_index(access_token)
+            if index < 0:
+                return None
+            next_item = dict(self._accounts[index])
+            history = list(next_item.get("recent_img2_results") or [])
+            history.append(classification)
+            if recent_window > 0:
+                history = history[-recent_window:]
+            next_item["recent_img2_results"] = history
+            if classification == "gptimage1":
+                next_item["consecutive_img1_results"] = int(next_item.get("consecutive_img1_results") or 0) + 1
+            else:
+                next_item["consecutive_img1_results"] = 0
+
+            should_demote = (
+                "gpt-image-2" in (next_item.get("preferred_models") or [])
+                and (
+                    int(next_item.get("consecutive_img1_results") or 0) >= consecutive_img1_threshold
+                    or history.count("gptimage1") >= recent_img1_threshold
+                )
+            )
+            if should_demote:
+                next_item["preferred_models"] = ["gpt-image-1"]
+                next_item["routing_priority"] = 0
+                next_item["auto_managed"] = True
+                next_item["last_probe_result"] = "auto_demoted"
+                next_item["demotion_count"] = int(next_item.get("demotion_count") or 0) + 1
+
+            account = self._normalize_account(next_item)
+            if account is None:
+                return None
+            account["_auto_demoted"] = should_demote
+            self._accounts[index] = {k: v for k, v in account.items() if not str(k).startswith("_")}
+            self._save_accounts()
+            return dict(account)
+
     def fetch_remote_info(self, access_token: str) -> dict[str, Any]:
         access_token = self._clean_token(access_token)
         if not access_token: