litellm-config/generate_keys.py

import os
import yaml
import hashlib
from urllib.parse import urlparse  # ✅ 必须导入
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):
    """
    [同步更新] 生成 api_base 的 Hash，强制忽略 URL 的路径部分。
    保持与 daily_limit.py v20 逻辑完全一致。
    """
    if not api_base:
        return "default"

    try:
        # 1. 预处理：去空格
        s = api_base.strip()
        
        # 2. 解析 URL
        parsed = urlparse(s)
        
        # 3. 提取 Scheme 和 Netloc (域名+端口)
        # 如果解析成功且包含 scheme 和 netloc，则重组
        if parsed.scheme and parsed.netloc:
            clean_base = f"{parsed.scheme}://{parsed.netloc}".lower()
        else:
            # 兜底：如果不是标准 URL，或者解析失败，尝试简单的去斜杠和小写化
            clean_base = s.rstrip('/').lower()
            
    except Exception:
        # 极端情况下的兜底
        clean_base = api_base.strip().rstrip('/').lower()

    # 4. 计算 Hash
    return hashlib.md5(clean_base.encode()).hexdigest()[:8]


def main():
    print("🔍 Generating Daily Limit Keys from config.yaml (Robust Mode)...\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
        # ⚠️ 这里非常重要：必须先解析出真实的 URL，再进行 Hash
        api_base_resolved = resolve_env_var(api_base_raw, env_vars)

        # Generate components
        clean_upstream = clean_model_name(upstream_model)
        
        # 使用同步后的 Hash 逻辑
        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()