Update sync_cliproxy_cleanup.py

sync_cliproxy_cleanup.py

@@ -11,24 +11,71 @@ LITELLM_MASTER_KEY = os.environ["LITELLM_MASTER_KEY"]
 CREDENTIAL_NAME = os.environ["LITELLM_CREDENTIAL_NAME"]          # LiteLLM 中预先创建好的凭证名称
 PRIMARY_MODEL_GROUP = os.environ.get("FALLBACK_PRIMARY_MODEL", "cliproxy/*")
 SYNC_INTERVAL = int(os.environ.get("SYNC_INTERVAL_SECONDS", 600))  # 默认10分钟
+MAX_RETRIES = int(os.environ.get("MAX_CREDENTIAL_RETRIES", 5))
+RETRY_DELAY = int(os.environ.get("RETRY_DELAY_SECONDS", 5))
 
 SYNC_TAG = "cliproxy-synced"
 MASTER_HEADERS = {"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}
 
-# ---------- 凭证获取 ----------
+# ---------- 辅助函数：等待 LiteLLM 完全就绪 ----------
+def wait_for_litellm_ready(timeout: int = 180):
+    """轮询 /health 直到 200，超时则报错"""
+    start = time.time()
+    print("⏳ 等待 LiteLLM Proxy 启动...")
+    while time.time() - start < timeout:
+        try:
+            resp = requests.get(f"{LITELLM_BASE_URL}/health", headers=MASTER_HEADERS, timeout=5)
+            if resp.status_code == 200:
+                print("✅ LiteLLM Proxy 已就绪")
+                return
+        except requests.RequestException:
+            pass
+        time.sleep(5)
+    raise RuntimeError(f"❌ LiteLLM Proxy 未在 {timeout}s 内就绪，请检查服务状态")
+
+# ---------- 凭证获取（带重试）----------
 def get_cliproxy_credential() -> Optional[Dict[str, str]]:
-    resp = requests.get(f"{LITELLM_BASE_URL}/credentials", headers=MASTER_HEADERS)
-    resp.raise_for_status()
-    credentials = resp.json().get("data", [])
-    for cred in credentials:
-        if cred.get("credential_name") == CREDENTIAL_NAME:
-            cred_info = cred.get("credential_info", {})
-            api_base = cred_info.get("api_base", "")
-            api_key = cred_info.get("api_key", "")
-            if not api_base or not api_key:
-                raise ValueError(f"凭证 '{CREDENTIAL_NAME}' 中缺少 api_base 或 api_key")
-            return {"api_base": api_base, "api_key": api_key}
-    raise ValueError(f"未找到名为 '{CREDENTIAL_NAME}' 的凭证，请确认凭证已创建")
+    """获取指定凭证，支持重试，防止数据库尚未加载完成"""
+    last_exception = None
+    for attempt in range(1, MAX_RETRIES + 1):
+        try:
+            resp = requests.get(f"{LITELLM_BASE_URL}/credentials", headers=MASTER_HEADERS)
+            resp.raise_for_status()
+            credentials = resp.json().get("data", [])
+            # 兼容有时直接返回数组的情况
+            if not isinstance(credentials, list):
+                credentials = resp.json() if isinstance(resp.json(), list) else []
+
+            for cred in credentials:
+                if cred.get("credential_name") == CREDENTIAL_NAME:
+                    cred_info = cred.get("credential_info", {})
+                    api_base = cred_info.get("api_base", "")
+                    api_key = cred_info.get("api_key", "")
+                    if not api_base or not api_key:
+                        raise ValueError(f"凭证 '{CREDENTIAL_NAME}' 中缺少 api_base 或 api_key")
+                    print(f"🔑 成功获取凭证 '{CREDENTIAL_NAME}'")
+                    return {"api_base": api_base, "api_key": api_key}
+
+            print(f"⚠️ 第 {attempt}/{MAX_RETRIES} 次查找凭证 '{CREDENTIAL_NAME}' 未找到，等待 {RETRY_DELAY}s 重试...")
+            last_exception = None
+            time.sleep(RETRY_DELAY)
+
+        except requests.RequestException as e:
+            print(f"⚠️ 第 {attempt}/{MAX_RETRIES} 次请求凭证列表失败: {e}")
+            last_exception = e
+            time.sleep(RETRY_DELAY)
+        except ValueError as e:
+            # 凭证结构错误，直接抛出
+            raise e
+
+    # 所有重试失败
+    if last_exception:
+        raise RuntimeError(f"❌ 无法获取凭证列表，网络或服务异常") from last_exception
+    raise ValueError(
+        f"❌ 未找到名为 '{CREDENTIAL_NAME}' 的凭证。\n"
+        f"   请确认凭证已在 UI 中创建，且 Secret 'LITELLM_SALT_KEY' 已设置为固定值。\n"
+        f"   如仍有问题，请连接数据库执行清理：DELETE FROM \"LiteLLM_VerificationToken\"; 后重启 LiteLLM。"
+    )
 
 # ---------- CLIProxy 模型获取 ----------
 def get_cliproxy_models(cred: dict) -> List[Dict]:
@@ -64,7 +111,7 @@ def delete_model(model_name: str):
     )
     return resp.status_code
 
-def add_model_to_litellm(model_data: Dict):
+def add_model_to_litellm(model_data: Dict) -> (bool, dict):
     original_id = model_data["id"]
     owner = model_data.get("owned_by", "unknown")
     new_name = f"{owner}/{original_id}" if not original_id.startswith(f"{owner}/") else original_id
@@ -153,6 +200,8 @@ def sync():
 
 # ---------- 守护进程入口 ----------
 if __name__ == "__main__":
+    wait_for_litellm_ready()          # 启动前等待 LiteLLM 完全就绪
+    print("开始守护同步任务...")
     while True:
         sync()
         time.sleep(SYNC_INTERVAL)