Update sync_cliproxy_cleanup.py

sync_cliproxy_cleanup.py

@@ -1,159 +1,306 @@
+`python
 import os
 import time
 import requests
 from typing import List, Dict, Set
 
-# ---------- 环境变量 ----------
-LITELLM_BASE_URL = os.environ.get("LITELLM_BASE_URL", "http://localhost:7860")
-LITELLM_MASTER_KEY = os.environ["LITELLM_MASTER_KEY"]
-CLIPROXY_BASE_URL = os.environ["CLIPROXY_BASE_URL"]
-CLIPROXY_API_KEY = os.environ.get("CLIPROXY_API_KEY", "")
-CREDENTIAL_NAME = os.environ.get("LITELLM_CREDENTIAL_NAME", "cliproxy")
-FORCE_CREDENTIAL = os.environ.get("FORCE_USE_CREDENTIAL", "false").lower() == "true"
-PRIMARY_MODEL_GROUP = os.environ.get("FALLBACK_PRIMARY_MODEL", "cliproxy/*")
-SYNC_INTERVAL = int(os.environ.get("SYNC_INTERVAL_SECONDS", 3600))
+---------- 环境变量 ----------
+LITELLMBASEURL = os.environ.get("LITELLMBASEURL", "http://localhost:7860")
+LITELLMMASTERKEY = os.environ["LITELLMMASTERKEY"]
+CLIPROXYBASEURL = os.environ["CLIPROXYBASEURL"]
+CLIPROXYAPIKEY = os.environ.get("CLIPROXYAPIKEY", "")
+CREDENTIALNAME = os.environ.get("LITELLMCREDENTIAL_NAME", "cliproxy")
+FORCECREDENTIAL = os.environ.get("FORCEUSE_CREDENTIAL", "false").lower() == "true"
+PRIMARYMODELGROUP = os.environ.get("FALLBACKPRIMARYMODEL", "cliproxy/*")
+SYNCINTERVAL = int(os.environ.get("SYNCINTERVAL_SECONDS", 3600))
 
 SYNC_TAG = "cliproxy-synced"
-MASTER_HEADERS = {"Authorization": f"Bearer {LITELLM_MASTER_KEY}"}
+MASTERHEADERS = {"Authorization": f"Bearer {LITELLMMASTER_KEY}"}
 
-# ---------- provider 推断 ----------
-def infer_provider(owner: str, model_id: str) -> str:
-    owner = owner.lower()
-    model_id = model_id.lower()
+provider 优先级（用于 fallback 排序）
+PROVIDER_PRIORITY = {
+    "cliproxy": 0,
+    "moonshotai": 1,
+    "kimi": 1,
+    "anthropic": 2,
+    "google": 3,
+}
 
-    if "gemini" in model_id or owner == "google":
-        return "gemini"
-    if "claude" in model_id or owner == "anthropic":
-        return "anthropic"
-    if owner in ("moonshotai", "kimi"):
-        return "openai"  # Moonshot 兼容 OpenAI 协议
-    if owner == "cliproxy":
-        return "openai"  # Cliproxy 本身是 OpenAI-compatible
-    return "openai"
 
-# ---------- 等待 LiteLLM ----------
-def wait_for_litellm_ready(timeout: int = 180):
+---------- 等待 LiteLLM 就绪 ----------
+def waitforlitellm_ready(timeout: int = 180):
     print("⏳ 等待 LiteLLM Proxy 启动...")
     start = time.time()
     while time.time() - start < timeout:
         try:
-            r = requests.get(f"{LITELLM_BASE_URL}/health", headers=MASTER_HEADERS, timeout=5)
-            if r.status_code == 200:
+            resp = requests.get(f"{LITELLMBASEURL}/health", headers=MASTER_HEADERS, timeout=5)
+            if resp.status_code == 200:
                 print("✅ LiteLLM Proxy 已就绪")
                 return
-        except Exception:
+        except requests.RequestException:
             pass
         time.sleep(5)
-    raise RuntimeError("❌ LiteLLM Proxy 启动超时")
+    raise RuntimeError(f"❌ LiteLLM Proxy 未在 {timeout}s 内就绪")
+
 
-# ---------- 凭证检查 ----------
-def credential_exists(name: str) -> bool:
-    r = requests.get(f"{LITELLM_BASE_URL}/credentials", headers=MASTER_HEADERS)
-    if not r.ok:
+---------- 检查凭证是否可用 ----------
+def credentialexists(credentialname: str) -> bool:
+    try:
+        resp = requests.get(f"{LITELLMBASEURL}/credentials", headers=MASTER_HEADERS)
+        if resp.status_code != 200:
+            print(f"⚠️ 获取凭证列表失败，HTTP {resp.status_code}: {resp.text}")
+            return False
+
+        data = resp.json()
+        if isinstance(data, list):
+            credentials = data
+        elif isinstance(data, dict) and "data" in data:
+            credentials = data["data"]
+        else:
+            print(f"⚠️ 无法识别的凭证列表格式: {data}")
+            return False
+
+        for cred in credentials:
+            if cred.get("credentialname") == credentialname or cred.get("name") == credential_name:
+                print(f"✅ 凭证 '{credential_name}' 存在")
+                return True
+        print(f"⚠️ 凭证 '{credential_name}' 不在列表中")
+        return False
+    except Exception as e:
+        print(f"⚠️ 检查凭证时出错: {e}")
         return False
-    data = r.json()
-    creds = data["data"] if isinstance(data, dict) else data
-    return any(c.get("credential_name") == name or c.get("name") == name for c in creds)
-
-# ---------- Cliproxy 模型 ----------
-def get_cliproxy_models() -> List[Dict]:
-    headers = {}
-    if CLIPROXY_API_KEY:
-        headers["Authorization"] = f"Bearer {CLIPROXY_API_KEY}"
-    r = requests.get(f"{CLIPROXY_BASE_URL}/models", headers=headers)
-    r.raise_for_status()
-    return r.json()["data"]
-
-# ---------- LiteLLM 模型 ----------
-def get_existing_models() -> List[Dict]:
-    r = requests.get(f"{LITELLM_BASE_URL}/v1/models", headers=MASTER_HEADERS)
-    r.raise_for_status()
-    return r.json().get("data", [])
-
-# ---------- 删除模型 ----------
-def delete_model(model_name: str):
-    r = requests.post(
-        f"{LITELLM_BASE_URL}/model/delete",
-        json={"model_name": model_name},
+
+
+---------- CLIProxy 模型获取 ----------
+def getcliproxymodels() -> List = {}
+    if CLIPROXYAPIKEY:
+        headers["Authorization"] = f"Bearer {CLIPROXYAPIKEY}"
+    resp = requests.get(f"{CLIPROXYBASEURL}/models", headers=headers)
+    resp.raiseforstatus()
+    return resp.json()["data"]
+
+
+---------- LiteLLM 现有模型 ----------
+def getexistingmodels_full() -> List[Dict]:
+    resp = requests.get(f"{LITELLMBASEURL}/v1/models", headers=MASTER_HEADERS)
+    resp.raiseforstatus()
+    return resp.json().get("data", [])
+
+
+---------- 模型删除 ----------
+def deletemodel(modelname: str) -> bool:
+    resp = requests.post(
+        f"{LITELLMBASEURL}/model/delete",
+        json={"modelname": modelname},
         headers=MASTER_HEADERS,
     )
-    return r.status_code in (200, 204)
+    if resp.status_code in [200, 204]:
+        print(f"  🗑️ 成功删除模型: {model_name}")
+        return True
+    else:
+        print(f"  ❌ 删除模型 {modelname} 失败，状态码: {resp.statuscode}, 响应: {resp.text}")
+        return False
+
+
+---------- 判断是否为同步模型（通过标签或 owner）----------
+def issyncedmodel(modelid: str, modelsfull: List[Dict]) -> bool:
+    for model in models_full:
+        if model["id"] == model_id:
+            if "tags" in model and SYNC_TAG in model["tags"]:
+                return True
+            owner = model.get("owned_by", "unknown")
+            if "/" in model_id and owner != "openai":
+                return True
+    return False
+
+
+---------- 推断 provider ----------
+def inferprovider(owner: str, modelid: str) -> str:
+    owner = owner.lower()
+    modelid = modelid.lower()
+
+    # cliproxy 永远走 OpenAI 兼容协议
+    if owner == "cliproxy":
+        return "openai"
+
+    if "gemini" in model_id or owner == "google":
+        return "gemini"
+    if "claude" in model_id or owner == "anthropic":
+        return "anthropic"
+    if owner in ("moonshotai", "kimi"):
+        return "openai"
+
+    return "openai"
+
+
+---------- 系统模型过滤（防止 openai/container 之类混入 fallback） ----------
+def issystemmodel(model_name: str) -> bool:
+    name = model_name.lower()
+    return (
+        name.startswith("container")
+        or name.startswith("hf")
+        or name.startswith("litellm")
+        or name.startswith("internal")
+    )
+
+
+---------- fallback 排序 ----------
+def sortfallbackmodels(models: Dict[str, str]) -> list[str]:
+    """
+    models: {model_name: owner}
+    """
+    def score(item):
+        model_name, owner = item
+        owner = owner.lower()
+
+        # embedding 永远排除（后面还会再过滤一遍）
+        if "embed" in model_name.lower():
+            return (99, model_name)
+
+        priority = PROVIDER_PRIORITY.get(owner, 50)
+        return (priority, model_name)
 
-# ---------- 添加模型（已修复） ----------
-def add_model_to_litellm(original_id: str, owner: str, use_credential: bool) -> bool:
-    provider = infer_provider(owner, original_id)
-    model_name = f"{owner}/{original_id}"
+    return [name for name, _ in sorted(models.items(), key=score)]
+
+
+---------- 模型添加 ----------
+def addmodeltolitellm(originalid: str, owner: str, use_credential: bool) -> bool:
+    provider = inferprovider(owner, originalid)
+    newname = f"{owner}/{originalid}" if not originalid.startswith(f"{owner}/") else originalid
 
     litellm_params = {
         "model": f"{provider}/{original_id}",
     }
 
-    # Gemini embedding 不支持 max_tokens
-    if provider == "gemini" and "embed" in original_id.lower():
-        pass
-
+    # Cliproxy / OpenAI 兼容：使用自定义 api_base + key
     if use_credential:
-        litellm_params["litellm_credential_name"] = CREDENTIAL_NAME
+        litellmparams["litellmcredentialname"] = CREDENTIALNAME
+        print(f"  🔑 新模型 {newname} 将引用凭证 '{CREDENTIALNAME}'")
     else:
-        litellm_params["api_base"] = CLIPROXY_BASE_URL.rstrip("/") + "/v1"
-        litellm_params["api_key"] = CLIPROXY_API_KEY
+        litellmparams["apibase"] = CLIPROXYBASEURL.rstrip("/") + "/v1"
+        litellmparams["apikey"] = CLIPROXYAPIKEY
+        print(f"  🔑 新模型 {new_name} 将使用环境变量中的 API Key (凭证回退)")
+
+    model_info = {
+        "owned_by": owner,
+        "tags": [SYNC_TAG],
+    }
+
+    # 标记模型类型，避免 LiteLLM 用 completion 去 probe embedding
+    if "embed" in original_id.lower():
+        modelinfo["modeltype"] = "embedding"
+    else:
+        modelinfo["modeltype"] = "chat"
 
     payload = {
-        "model_name": model_name,
-        "litellm_params": litellm_params,
-        "model_info": {
-            "owned_by": owner,
-            "tags": [SYNC_TAG],
-        },
+        "modelname": newname,
+        "litellmparams": litellmparams,
+        "modelinfo": modelinfo,
     }
 
-    r = requests.post(f"{LITELLM_BASE_URL}/model/new", json=payload, headers=MASTER_HEADERS)
-    if r.status_code in (200, 201):
-        print(f"  ➕ 新增模型 {model_name}")
+    resp = requests.post(
+        f"{LITELLMBASEURL}/model/new",
+        json=payload,
+        headers=MASTER_HEADERS,
+    )
+    if resp.status_code in (200, 201):
+        print(f"  ➕ 新增: {new_name}")
         return True
+    else:
+        print(f"  ❌ 添加模型 {newname} 失败: {resp.statuscode} {resp.text}")
+        return False
 
-    print(f"  ❌ 添加失败 {model_name}: {r.status_code} {r.text}")
-    return False
 
-# ---------- 同步 ----------
+---------- 更新 Fallback 链 ----------
+def updatefallbackchain(model_names: List[str]):
+    # 先清空旧规则
+    requests.post(
+        f"{LITELLMBASEURL}/fallback/delete",
+        json={"model": PRIMARYMODELGROUP, "fallback_type": "general"},
+        headers=MASTER_HEADERS,
+    )
+
+    if not model_names:
+        print("⚠️ 没有可用的 Fallback 模型，已清空规则。")
+        return
+
+    fallbackmodels = modelnames[:50]
+    payload = {
+        "model": PRIMARYMODELGROUP,
+        "fallbackmodels": fallbackmodels,
+        "fallback_type": "general",
+    }
+    resp = requests.post(f"{LITELLMBASEURL}/fallback", json=payload, headers=MASTER_HEADERS)
+    if resp.ok:
+        print(f"✅ Fallback 已更新，包含 {len(fallback_models)} 个模型。")
+    else:
+        print(f"❌ Fallback 更新失败: {resp.status_code} {resp.text}")
+
+
+---------- 主同步流程 ----------
 def sync():
-    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 开始同步")
+    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 开始同步…")
 
-    use_credential = FORCE_CREDENTIAL or credential_exists(CREDENTIAL_NAME)
+    usecredential = FORCECREDENTIAL or credentialexists(CREDENTIALNAME)
+    if not use_credential:
+        print(f"⚠️ 凭证 '{CREDENTIAL_NAME}' 未找到，将使用环境变量中的 API Key")
 
-    clip_models = get_cliproxy_models()
-    clip_map = {f"{m['owned_by']}/{m['id']}": m["owned_by"] for m in clip_models}
+    # 1. 获取 Cliproxy 模型
+    cliproxymodels = getcliproxy_models()
+    cliproxy_map: Dict[str, str] = {}
+    for m in cliproxy_models:
+        owner = m.get("owned_by", "cliproxy")
+        new_name = f"{owner}/{m['id']}"
+        cliproxymap[newname] = owner
 
-    existing = get_existing_models()
-    existing_ids = {m["id"] for m in existing}
+    # 2. 获取 LiteLLM 现有模型
+    existingfull = getexistingmodelsfull()
+    existingids: Set[str] = {m["id"] for m in existingfull}
 
-    for full_name, owner in clip_map.items():
-        if full_name not in existing_ids:
-            add_model_to_litellm(full_name.split("/", 1)[1], owner, use_credential)
+    # 3. 新增 / 迁移模型
+    added = 0
+    migrated = 0
+    for newname, owner in cliproxymap.items():
+        if newname not in existingids:
+            if addmodeltolitellm(newname.split("/", 1)[1], owner, use_credential):
+                added += 1
+        else:
+            # 如果需要凭证，则通过“删除+重新添加”的方式迁移
+            if use_credential:
+                originalid = newname.split("/", 1)[1]
+                if deletemodel(newname):
+                    if addmodeltolitellm(originalid, owner, use_credential=True):
+                        print(f"  🔄 模型 {new_name} 已迁移为凭证模式")
+                        migrated += 1
+                    else:
+                        print(f"  ❌ 迁移 {new_name} 失败：重新添加失败")
+                else:
+                    print(f"  ⚠️ 无法删除模型 {new_name}，跳过迁移")
 
-    # 删除失效模型
-    for m in existing:
-        if SYNC_TAG in m.get("tags", []) and m["id"] not in clip_map:
-            delete_model(m["id"])
+    # 4. 删除失效模型（通过标签识别）
+    deleted = 0
+    for modelfull in existingfull:
+        modelid = modelfull["id"]
+        if issyncedmodel(modelid, existingfull) and modelid not in cliproxymap:
+            if deletemodel(modelid):
+                deleted += 1
 
-    # fallback（排除 embedding）
-    fallback_models = [m for m in clip_map if "embed" not in m.lower()]
-    requests.post(
-        f"{LITELLM_BASE_URL}/fallback",
-        json={
-            "model": PRIMARY_MODEL_GROUP,
-            "fallback_models": fallback_models[:50],
-            "fallback_type": "general",
-        },
-        headers=MASTER_HEADERS,
-    )
+    # 5. 生成 fallback 列表：排序 + 过滤 embedding + 过滤系统模型
+    sortedmodels = sortfallbackmodels(cliproxymap)
+    fallback_candidates = [
+        m for m in sorted_models
+        if "embed" not in m.lower() and not issystemmodel(m)
+    ]
+    updatefallbackchain(fallback_candidates)
+
+    print(f"✅ 同步完成：新增 {added}, 迁移凭证 {migrated}, 删除 {deleted}\n")
 
-    print("✅ 同步完成\n")
 
-# ---------- 守护 ----------
-if __name__ == "__main__":
-    wait_for_litellm_ready()
-    print(f"🚀 同步守护启动，每 {SYNC_INTERVAL}s 执行一次")
+---------- 守护���程 ----------
+if name == "main":
+    waitforlitellm_ready()
+    print("🚀 守护同步任务启动，间隔 " + str(SYNC_INTERVAL) + " 秒")
     while True:
         sync()
-        time.sleep(SYNC_INTERVAL)
+        time.sleep(SYNC_INTERVAL)
+`