src/main.py

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:
        # 从环境变量加载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_manager.add_key(
                key_value="Bearer eyJhbGciOiJSUzI1NiIsImtpZCI6IjE5MzQ0ZTY1LWJiYzktNDRkMS1hOWQwLWY5NTdiMDc5YmQwZSIsInR5cCI6IkpXVCJ9...",
                name="默认测试密钥"
            )
        
        logger.info(f"API服务初始化完成，已加载 {len(key_manager.keys)} 个API key")
    except Exception as e:
        logger.error(f"API服务初始化失败: {str(e)}")
        raise

def start():
    """启动API服务"""
    # 初始化应用
    init_app()
    
    # 打印配置信息
    Config.print_config()
    
    # 启动服务
    logger.info(f"启动服务: {Config.HOST}:{Config.PORT}")
    uvicorn.run(
        "src.app:app",
        host=Config.HOST,
        port=Config.PORT,
        reload=Config.VERBOSE_LOGGING  # 仅在调试模式下开启自动重载
    )

if __name__ == "__main__":
    start()