Update src/main.py

src/main.py

@@ -1,18 +1,85 @@
 import uvicorn
 import logging
+import json
+import os
 from .app import app, key_manager
 from .config import Config
 
 # 获取日志记录器
 logger = logging.getLogger("sora-api.main")
 
+def ensure_bearer_prefix(key):
+    """确保API密钥包含Bearer前缀"""
+    if key and not key.startswith("Bearer "):
+        return f"Bearer {key}"
+    return key
+
+def format_api_keys_json(keys):
+    """格式化API密钥为JSON字符串，确保包含Bearer前缀"""
+    if isinstance(keys, list):
+        for i, key_info in enumerate(keys):
+            if isinstance(key_info, dict) and "key" in key_info:
+                keys[i]["key"] = ensure_bearer_prefix(key_info["key"])
+    return json.dumps(keys)
+
 def init_app():
     """初始化应用程序"""
     try:
-        # 密钥管理器已在app.py中初始化并加载完成
+        # 从环境变量加载API密钥
+        api_keys_loaded = False
+        
+        # 从API_AUTH_TOKEN环境变量加载密钥
+        api_auth_token = os.getenv("API_AUTH_TOKEN", "")
+        if api_auth_token:
+            logger.info(f"从环境变量API_AUTH_TOKEN读取，长度: {len(api_auth_token)}")
+            
+            # 尝试解析为JSON（多个密钥）
+            try:
+                api_keys_data = json.loads(api_auth_token)
+                logger.info(f"API_AUTH_TOKEN包含JSON数据，尝试解析为密钥列表")
+                
+                # 清空现有密钥并设置新密钥
+                key_manager.keys = []
+                key_manager.usage_stats = {}
+                
+                # 如果是列表，作为多个密钥处理
+                if isinstance(api_keys_data, list):
+                    key_manager._process_keys_data(api_keys_data)
+                    logger.info(f"已从API_AUTH_TOKEN加载 {len(key_manager.keys)} 个API密钥")
+                    api_keys_loaded = True
+                # 如果是字典，检查是否包含keys字段
+                elif isinstance(api_keys_data, dict) and "keys" in api_keys_data:
+                    key_manager._process_keys_data(api_keys_data["keys"])
+                    logger.info(f"已从API_AUTH_TOKEN加载 {len(key_manager.keys)} 个API密钥")
+                    api_keys_loaded = True
+                else:
+                    logger.warning(f"API_AUTH_TOKEN包含JSON数据，但格式不是密钥列表")
+            except json.JSONDecodeError:
+                # 不是JSON，作为单个密钥处理
+                logger.info(f"API_AUTH_TOKEN不是JSON格式，作为单个密钥处理")
+                
+                # 清空现有密钥并设置新密钥
+                key_manager.keys = []
+                key_manager.usage_stats = {}
+                
+                # 添加从环境变量读取的单个密钥
+                key_manager.add_key(
+                    key_value=ensure_bearer_prefix(api_auth_token),
+                    name="环境变量密钥",
+                    weight=1,
+                    rate_limit=60
+                )
+                
+                logger.info(f"已将API_AUTH_TOKEN作为单个密钥设置")
+                api_keys_loaded = True
+        
+        # 保存密钥到文件
+        if api_keys_loaded:
+            key_manager._save_keys()
+        
         # 检查是否有可用的密钥
         if not key_manager.keys:
-            logger.warning("未配置API key，将使用测试密钥")
+            logger.warning("未从环境变量加载到API密钥，将使用测试密钥")
             key_manager.add_key(
                 key_value="Bearer eyJhbGciOiJSUzI1NiIsImtpZCI6IjE5MzQ0ZTY1LWJiYzktNDRkMS1hOWQwLWY5NTdiMDc5YmQwZSIsInR5cCI6IkpXVCJ9...",
                 name="默认测试密钥"