初始化空白分支文件

generate_keys.py

@@ -0,0 +1,153 @@
+import os
+import yaml
+import hashlib
+import re
+from typing import Dict, Any
+
+# ==========================================
+# 🛠️ 配置
+# ==========================================
+CONFIG_FILE = "config.yaml"
+ENV_FILE = ".env"
+
+
+def load_env(env_path: str) -> Dict[str, str]:
+    """读取 .env 文件并返回字典"""
+    env_vars = {}
+    if not os.path.exists(env_path):
+        print(f"⚠️ Warning: {env_path} not found. Env var substitution might fail.")
+        return env_vars
+
+    with open(env_path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            try:
+                key, value = line.split("=", 1)
+                env_vars[key.strip()] = value.strip().strip('"').strip("'")
+            except ValueError:
+                continue
+    return env_vars
+
+
+def resolve_env_var(value: str, env_vars: Dict[str, str]) -> str:
+    """
+    处理环境变量替换，支持以下两种格式：
+    1. ${VAR_NAME} - 标准格式
+    2. os.environ/VAR_NAME - LiteLLM 特有格式
+    """
+    if not isinstance(value, str):
+        return value
+
+    # 1. 处理 os.environ/VAR_NAME 格式
+    if value.startswith("os.environ/"):
+        var_name = value.replace("os.environ/", "")
+        return env_vars.get(var_name, os.getenv(var_name, value))
+
+    # 2. 处理 ${VAR_NAME} 或 $VAR_NAME 格式
+    pattern = re.compile(r"\$\{?([A-Z0-9_]+)\}?")
+
+    def replace_match(match):
+        var_name = match.group(1)
+        # 优先从 .env 读取，其次从系统环境变量读取
+        return env_vars.get(var_name, os.getenv(var_name, match.group(0)))
+
+    return pattern.sub(replace_match, value)
+
+
+def clean_model_name(name):
+    if not name:
+        return ""
+    return name.replace("openai/", "").replace("azure/", "").replace("anthropic/", "")
+
+
+def generate_api_base_hash(api_base):
+    if not api_base:
+        return "default"
+    # Logic matches daily_limit.py v18
+    
+    # === 新增标准化代码 ===
+    # 1. 去除首尾空格 (strip)
+    # 2. 去除末尾的斜杠 (rstrip('/'))
+    # 这样 '.../v3/' 和 '.../v3' 会生成相同的 Hash
+    normalized_base = api_base.strip().rstrip('/')
+    # ====================
+
+    
+    return hashlib.md5(api_base.encode()).hexdigest()[:8]
+
+
+def main():
+    print("🔍 Generating Daily Limit Keys from config.yaml...\n")
+
+    # 1. Load Env
+    env_vars = load_env(ENV_FILE)
+
+    # 2. Load Config
+    if not os.path.exists(CONFIG_FILE):
+        print(f"❌ Error: {CONFIG_FILE} not found!")
+        return
+
+    try:
+        with open(CONFIG_FILE, "r", encoding="utf-8") as f:
+            config = yaml.safe_load(f)
+    except Exception as e:
+        print(f"❌ Error parsing yaml: {e}")
+        return
+
+    model_list = config.get("model_list", [])
+    if not model_list:
+        print("⚠️ No models found in config.yaml")
+        return
+
+    results = []
+    seen_keys = set()
+
+    # Table Header
+    # Model Alias | Upstream Model | API Base (Env/URL) | Generated Key
+    header = f"{'Model Name (Alias)':<25} | {'Upstream Model':<30} | {'API Base (Env/URL)':<40} | {'Generated Key'}"
+    print(header)
+    print("-" * len(header))
+
+    for model in model_list:
+        litellm_params = model.get("litellm_params", {})
+
+        # Get raw values
+        model_alias = model.get("model_name", "unknown")
+        upstream_model = litellm_params.get("model", "")
+        api_base_raw = litellm_params.get("api_base", "")
+
+        # Resolve env vars for hash generation
+        api_base_resolved = resolve_env_var(api_base_raw, env_vars)
+
+        # Generate components
+        clean_upstream = clean_model_name(upstream_model)
+        base_hash = generate_api_base_hash(api_base_resolved)
+
+        # Generate final key
+        final_key = f"{clean_upstream}:{base_hash}"
+
+        # Deduplicate
+        if final_key in seen_keys:
+            continue
+        seen_keys.add(final_key)
+
+        # Truncate api_base_raw if it's too long for display
+        display_api_base = api_base_raw
+        if len(display_api_base) > 37:
+            display_api_base = display_api_base[:34] + "..."
+
+        print(f"{model_alias:<25} | {clean_upstream:<30} | {display_api_base:<40} | {final_key}")
+        results.append(f'"{final_key}": 1500000')
+
+    print("\n📋 [Copy & Paste into daily_limit.py MODEL_QUOTAS]:")
+    print("-" * 50)
+    print("MODEL_QUOTAS = {")
+    print(",\n".join(f"    {item}" for item in results))
+    print("}")
+    print("-" * 50)
+
+
+if __name__ == "__main__":
+    main()