Update api_service.py

api_service.py

@@ -1,339 +1,309 @@
-#!/usr/bin/env python3
 """
-Warp API 服务器
-整合所有模块，提供OpenAI兼容的API接口和Token管理Web界面
+API服务模块
+处理所有API请求逻辑，包括聊天完成和Token管理
 """
 import json
-from flask import Flask, request, jsonify, Response, stream_with_context, render_template_string, session, redirect, url_for
+import time
+from datetime import datetime
+from typing import List, Dict, Any, Generator
 from loguru import logger
 
 from config import Config
 from utils import Utils
-from api_service import ApiService
+from token_manager import MultiTokenManager
+from warp_client import WarpClient
+from request_converter import RequestConverter
+from model_mapper import ModelMapper
 
-# 初始化日志
-Utils.setup_logging()
 
-# 创建Flask应用
-app = Flask(__name__)
-app.secret_key = 'warp-api-server-secret-key-2024'  # 用于session管理
-
-# 创建API服务
-api_service = ApiService()
-
-def check_admin_auth():
-    """检查管理员认证"""
-    if not Config.require_admin_auth():
-        return True
-    
-    # 检查session中的管理员认证状态
-    return session.get('admin_authenticated', False)
-
-def verify_admin_key(admin_key: str) -> bool:
-    """验证管理员密钥"""
-    if not Config.require_admin_auth():
-        return True
-    
-    return admin_key == Config.get_admin_key()
-
-@app.before_request
-def check_auth():
-    """请求前鉴权检查"""
-    # 跳过管理员认证的端点
-    admin_auth_endpoints = ['admin_login', 'admin_auth']
-    if request.endpoint in admin_auth_endpoints:
-        return
-    
-    # 管理界面相关端点需要管理员认证
-    admin_endpoints = ['index', 'token_status', 'add_tokens', 'batch_login', 'refresh_tokens', 'remove_tokens', 'export_tokens']
-    if request.endpoint in admin_endpoints:
-        if not check_admin_auth():
-            if request.endpoint == 'index':
-                return redirect(url_for('admin_login'))
-            else:
-                return jsonify({"error": "需要管理员认证"}), 401
-    
-    # API端点需要API密钥认证
-    api_endpoints = ['get_models', 'chat_completions']
-    if request.endpoint in api_endpoints:
-        auth_header = request.headers.get('Authorization')
-        if not api_service.authenticate_request(auth_header):
-            return jsonify({"error": "未授权访问"}), 401
-
-@app.route('/admin/login')
-def admin_login():
-    """管理员登录页面"""
-    if not Config.require_admin_auth():
-        return redirect(url_for('index'))
-    
-    if check_admin_auth():
-        return redirect(url_for('index'))
+class ApiService:
+    """API服务类，处理所有业务逻辑"""
     
-    from web_template import get_admin_login_template
-    return render_template_string(get_admin_login_template())
-
-@app.route('/admin/auth', methods=['POST'])
-def admin_auth():
-    """管理员认证处理"""
-    if not Config.require_admin_auth():
-        return jsonify({"success": True, "redirect": "/"})
-    
-    try:
-        data = request.get_json()
-        admin_key = data.get('admin_key', '')
+    def __init__(self):
+        # 初始化Token管理器
+        self.token_manager = MultiTokenManager()
         
-        if verify_admin_key(admin_key):
-            session['admin_authenticated'] = True
-            logger.info("🔐 管理员认证成功")
-            return jsonify({"success": True, "message": "认证成功", "redirect": "/"})
-        else:
-            logger.warning("⚠️ 管理员认证失败")
-            return jsonify({"success": False, "message": "管理员密钥错误"})
-            
-    except Exception as e:
-        logger.error(f"❌ 管理员认证出错: {e}")
-        return jsonify({"success": False, "message": "认证失败"}), 500
-
-@app.route('/admin/logout', methods=['POST'])
-def admin_logout():
-    """管理员登出"""
-    session.pop('admin_authenticated', None)
-    return jsonify({"success": True, "message": "已登出"})
-
-@app.route('/')
-def index():
-    """管理页面"""
-    from web_template import get_html_template
-    return render_template_string(get_html_template())
-
-@app.route(f"/{Config.OPENAI_API_VERSION}/models", methods=['GET'])
-def get_models():
-    """获取模型列表"""
-    return jsonify(api_service.get_models())
-
-@app.route(f"/{Config.OPENAI_API_VERSION}/chat/completions", methods=['POST'])
-def chat_completions():
-    """聊天完成端点"""
-    try:
-        request_data = request.get_json()
-        if not request_data:
-            return jsonify({"error": "无效的请求数据"}), 400
+        # 初始化Warp客户端
+        self.warp_client = WarpClient(self.token_manager)
         
-        # 检查模型是否有效
-        model = request_data.get("model")
-        from model_mapper import ModelMapper
-        if not ModelMapper.is_valid_model(model):
-            return jsonify({"error": f"不支持的模型: {model}"}), 400
-        
-        stream = request_data.get("stream", False)
-        
-        if stream:
-            # 流式响应
-            def generate():
-                for chunk in api_service.chat_completion(request_data, stream=True):
-                    yield chunk
+        logger.info("🚀 ApiService初始化完成")
+    
+    def authenticate_request(self, auth_header: str) -> bool:
+        """验证API请求"""
+        if not auth_header:
+            return False
             
-            return Response(
-                stream_with_context(generate()),
-                content_type='text/event-stream',
-                headers={
-                    'Cache-Control': 'no-cache',
-                    'Connection': 'keep-alive',
-                    'Access-Control-Allow-Origin': '*'
-                }
-            )
-        else:
-            # 非流式响应
-            response_text = ""
-            for chunk in api_service.chat_completion(request_data, stream=False):
-                response_text = chunk
-                break
+        token = Utils.extract_bearer_token(auth_header)
+        if not token:
+            return False
             
-            return Response(
-                json.dumps(response_text) if isinstance(response_text, dict) else response_text,
-                content_type='application/json'
-            )
+        return Utils.validate_api_key(token)
     
-    except Exception as e:
-        logger.error(f"❌ 处理聊天请求时出错: {e}")
-        return jsonify({"error": f"服务器内部错误: {str(e)}"}), 500
-
-@app.route('/token/status', methods=['GET'])
-def token_status():
-    """获取token状态"""
-    try:
-        status = api_service.get_token_status()
-        return jsonify(status)
-    except Exception as e:
-        return jsonify({"success": False, "message": str(e)}), 500
-
-@app.route('/token/add', methods=['POST'])
-def add_tokens():
-    """添加token"""
-    try:
-        data = request.get_json()
-        tokens = data.get('tokens', [])
-        
-        result = api_service.add_tokens(tokens)
-        return jsonify(result)
-        
-    except Exception as e:
-        return jsonify({"success": False, "message": str(e)}), 500
-
-@app.route('/token/remove', methods=['POST'])
-def remove_tokens():
-    """删除token"""
-    try:
-        data = request.get_json()
-        refresh_token = data.get('refresh_token')
+    def get_models(self) -> Dict[str, Any]:
+        """获取支持的模型列表"""
+        models = []
+        for model_name in ModelMapper.get_available_models():
+            models.append({
+                "id": model_name,
+                "object": "model",
+                "created": Utils.get_current_timestamp(),
+                "owned_by": "warp-proxy"
+            })
         
-        result = api_service.remove_refresh_token(refresh_token=refresh_token)
-        return jsonify(result)
-        
-    except Exception as e:
-        return jsonify({"success": False, "message": str(e)}), 500
-
-@app.route('/token/refresh', methods=['POST'])
-def refresh_tokens():
-    """刷新所有token"""
-    try:
-        result = api_service.refresh_all_tokens()
-        return jsonify(result)
-    except Exception as e:
-        return jsonify({"success": False, "message": str(e)}), 500
-
-@app.route('/token/export', methods=['POST'])
-def export_tokens():
-    """导出refresh token（需要超级管理员密钥验证）"""
-    try:
-        data = request.get_json()
-        super_admin_key = data.get('super_admin_key', '')
-        
-        result = api_service.export_refresh_tokens(super_admin_key)
-        return jsonify(result)
+        return {
+            "object": "list",
+            "data": models
+        }
+    
+    def chat_completion(self, request_data: Dict[str, Any], stream: bool = False) -> Generator[str, None, None]:
+        """处理聊天完成请求"""
+        request_id = Utils.generate_request_id()
+        # 解析请求
+        openai_request = RequestConverter.parse_openai_request(request_data)
+        model = ModelMapper.get_warp_model(openai_request.model)
+        messages = openai_request.messages
         
-    except Exception as e:
-        logger.error(f"❌ 导出token时出错: {e}")
-        return jsonify({"success": False, "message": str(e)}), 500
-
-@app.route('/login/batch', methods=['POST'])
-def batch_login():
-    """批量获取refresh token"""
-    try:
-        data = request.get_json()
-        email_url_pairs = data.get('email_url_pairs', [])
-        max_workers = data.get('max_workers', 5)
+        logger.info(f"🎯 开始处理聊天请求 [ID: {request_id[:8]}] [模型: {model}] [流式: {stream}]")
+        start_time = time.time()
         
-        # 转换为字典格式
-        email_url_dict = {}
-        for pair in email_url_pairs:
-            email = pair.get('email', '').strip()
-            url = pair.get('url', '').strip()
+        try:
+            # 创建protobuf数据
+            protobuf_data = self.warp_client.create_protobuf_data(messages, model)
+            if not protobuf_data:
+                error_msg = "创建请求数据失败"
+                logger.error(f"❌ {error_msg} [ID: {request_id[:8]}]")
+                yield self._create_error_response(error_msg, request_id)
+                return
             
-            if email and url:
-                email_url_dict[email] = url
-        
-        if not email_url_dict:
-            return jsonify({'success': False, 'message': '没有有效的邮箱和URL对'}), 400
-        
-        result = api_service.batch_get_refresh_tokens(email_url_dict, max_workers)
-        return jsonify(result)
-        
-    except Exception as e:
-        logger.error(f"❌ 批量登录时出错: {e}")
-        return jsonify({'success': False, 'message': f'服务器错误: {str(e)}'}), 500
-
-@app.route('/health', methods=['GET'])
-def health_check():
-    """健康检查端点"""
-    return jsonify({
-        "status": "healthy", 
-        "timestamp": Utils.get_current_timestamp(),
-        "version": "2.0.0"
-    })
-
-def compile_protobuf_at_startup():
-    """启动时编译protobuf"""
-    logger.info("🔧 开始编译Protobuf文件...")
-    try:
-        from protobuf_manager import ProtobufManager
-        success = ProtobufManager.validate_protobuf_module()
-        if success:
-            logger.success("✅ Protobuf模块验证成功")
-        else:
-            logger.warning("⚠️ Protobuf模块验证失败，但服务仍将启动")
-    except Exception as e:
-        logger.error(f"❌ 编译Protobuf时出错: {e}")
-
-def main():
-    """主函数"""
+            # 发送请求并处理响应
+            response_chunks = 0
+            total_content = ""
+            
+            logger.success(f"🚀 开始接收响应 [ID: {request_id[:8]}]")
+            
+            for chunk_text in self.warp_client.send_request(protobuf_data):
+                if chunk_text:
+                    response_chunks += 1
+                    total_content += chunk_text
+                    
+                    logger.debug(f"📦 响应块 #{response_chunks} [ID: {request_id[:8]}] [长度: {len(chunk_text)}]")
+                    
+                    if stream:
+                        # 流式响应
+                        chunk_response = self._create_stream_chunk(chunk_text, request_id)
+                        yield f"data: {json.dumps(chunk_response)}\n\n"
+                    else:
+                        # 非流式响应 - 等待完整内容
+                        continue
+            
+            # 处理响应结束
+            end_time = time.time()
+            duration = end_time - start_time
+            
+            if stream:
+                # 发送结束标记
+                final_chunk = self._create_stream_end_chunk(request_id)
+                yield f"data: {json.dumps(final_chunk)}\n\n"
+                yield "data: [DONE]\n\n"
+                
+                logger.success(f"✅ 流式响应完成 [ID: {request_id[:8]}] [块数: {response_chunks}] [耗时: {duration:.2f}s]")
+            else:
+                # 返回完整响应
+                response = self._create_complete_response(total_content, request_id)
+                yield response
+                
+                logger.success(f"✅ 完整响应完成 [ID: {request_id[:8]}] [长度: {len(total_content)}] [耗时: {duration:.2f}s]")
+                
+        except Exception as e:
+            logger.error(f"❌ 聊天请求处理失败 [ID: {request_id[:8]}]: {e}")
+            yield self._create_error_response(f"服务器内部错误: {str(e)}", request_id)
     
-    print("=" * 80)
-    print("🚀 启动 Warp API 服务器 v2.0")
-    print("=" * 80)
-    print(f"📡 服务器地址: http://localhost:{Config.SERVER_PORT}")
-    print(f"🔧 管理界面: http://localhost:{Config.SERVER_PORT}")
-    print(f"📚 模型端点: http://localhost:{Config.SERVER_PORT}/{Config.OPENAI_API_VERSION}/models")
-    print(f"💬 聊天端点: http://localhost:{Config.SERVER_PORT}/{Config.OPENAI_API_VERSION}/chat/completions")
-    print(f"❤️  健康检查: http://localhost:{Config.SERVER_PORT}/health")
-    print(f"📝 文件日志: {'启用' if Config.enable_file_logging() else '禁用'}")
+    def _create_stream_chunk(self, content: str, request_id: str) -> Dict[str, Any]:
+        """创建流式响应块"""
+        return {
+            "id": f"chatcmpl-{request_id}",
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": "gemini-2.0-flash",
+            "choices": [{
+                "index": 0,
+                "delta": {"content": content},
+                "finish_reason": None
+            }]
+        }
     
-    # 显示安全配置信息
-    if Config.require_admin_auth():
-        print(f"🔐 管理员认证: 已启用 (需要 ADMIN_KEY)")
-    else:
-        print("⚠️  管理员认证: 未启用 (可通过 ADMIN_KEY 环境变量启用)")
+    def _create_stream_end_chunk(self, request_id: str) -> Dict[str, Any]:
+        """创建流式响应结束块"""
+        return {
+            "id": f"chatcmpl-{request_id}",
+            "object": "chat.completion.chunk", 
+            "created": int(time.time()),
+            "model": "gemini-2.0-flash",
+            "choices": [{
+                "index": 0,
+                "delta": {},
+                "finish_reason": "stop"
+            }]
+        }
     
-    if Config.require_super_admin_auth():
-        print(f"🔒 超级管理员认证: 已启用 (导出功能需要 SUPER_ADMIN_KEY)")
-    else:
-        print("⚠️  超级管理员认证: 未启用 (可通过 SUPER_ADMIN_KEY 环境变量启用)")
+    def _create_complete_response(self, content: str, request_id: str) -> Dict[str, Any]:
+        """创建完整响应"""
+        return {
+            "id": f"chatcmpl-{request_id}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": "gemini-2.0-flash",
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": content
+                },
+                "finish_reason": "stop"
+            }],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            }
+        }
     
-    # 显示环境变量token信息
-    env_tokens = Config.get_refresh_tokens()
-    if env_tokens:
-        print(f"🎯 环境变量Token: 已设置 {len(env_tokens)} 个")
-        for i, token in enumerate(env_tokens, 1):
-            print(f"   #{i}: {token[:20]}...{token[-8:] if len(token) > 28 else ''}")
-    else:
-        print("⚠️  环境变量Token: 未设置 (请设置 WARP_REFRESH_TOKEN)")
-        print("💡 支持多个token，使用分号(;)分割：token1;token2;token3")
+    def _create_error_response(self, error_message: str, request_id: str) -> Dict[str, Any]:
+        """创建错误响应"""
+        return {
+            "error": {
+                "message": error_message,
+                "type": "api_error",
+                "code": "internal_error"
+            },
+            "id": request_id
+        }
     
-    print("-" * 80)
+    def get_token_status(self) -> Dict[str, Any]:
+        """获取Token状态"""
+        try:
+            status = self.token_manager.get_token_status()
+            return {"success": True, **status}
+        except Exception as e:
+            logger.error(f"❌ 获取Token状态失败: {e}")
+            return {"success": False, "message": str(e)}
     
-    try:
-        # 编译protobuf
-        compile_protobuf_at_startup()
-        
-        # 启动后台服务
-        logger.info("🔄 启动后台服务...")
-        api_service.start_services()
-        logger.success("✅ 后台服务启动完成")
-        
-        logger.info("-" * 80)
-        logger.success("🌟 服务器启动完成！")
-        if Config.require_admin_auth():
-            logger.info(f"🔐 访问管理界面需要管理员密钥: http://localhost:{Config.SERVER_PORT}")
-        else:
-            logger.info(f"🌐 访问管理界面: http://localhost:{Config.SERVER_PORT}")
-        logger.info("-" * 80)
-        
-        # 启动Flask服务器
-        app.run(
-            host=Config.SERVER_HOST,
-            port=Config.SERVER_PORT,
-            debug=False,
-            threaded=True
-        )
-        
-    except KeyboardInterrupt:
-        logger.info("🛑 接收到停止信号")
-    except Exception as e:
-        logger.error(f"❌ 服务器启动失败: {e}")
-    finally:
-        logger.info("🧹 清理资源...")
-        api_service.stop_services()
-        logger.info("👋 服务器已停止")
-
-if __name__ == "__main__":
-    main()
+    def add_tokens(self, tokens: List[str]) -> Dict[str, Any]:
+        """添加Token"""
+        try:
+            success = self.token_manager.add_refresh_tokens(tokens)
+            if success:
+                valid_tokens = [t for t in tokens if Utils.validate_refresh_token_format(t)]
+                return {
+                    "success": True, 
+                    "message": "Token添加成功", 
+                    "added_tokens": len(valid_tokens)
+                }
+            else:
+                return {"success": False, "message": "没有有效的Token可添加"}
+        except Exception as e:
+            logger.error(f"❌ 添加Token失败: {e}")
+            return {"success": False, "message": str(e)}
+    
+    def remove_refresh_token(self, refresh_token: str) -> Dict[str, Any]:
+        """删除refresh token"""
+        try:
+            success = self.token_manager.remove_refresh_token(refresh_token)
+            if success:
+                return {"success": True, "message": "Token删除成功"}
+            else:
+                return {"success": False, "message": "Token不存在"}
+        except Exception as e:
+            logger.error(f"❌ 删除Token失败: {e}")
+            return {"success": False, "message": str(e)}
+    
+    def refresh_all_tokens(self) -> Dict[str, Any]:
+        """刷新所有Token"""
+        try:
+            self.token_manager.refresh_all_tokens()
+            return {"success": True, "message": "Token刷新已开始"}
+        except Exception as e:
+            logger.error(f"❌ 刷新Token失败: {e}")
+            return {"success": False, "message": str(e)}
+    
+    def export_refresh_tokens(self, super_admin_key: str) -> Dict[str, Any]:
+        """导出refresh token内容（需要超级管理员密钥验证）"""
+        try:
+            # 验证超级管理员密钥
+            if Config.require_super_admin_auth():
+                if not super_admin_key or super_admin_key != Config.get_super_admin_key():
+                    return {"success": False, "message": "超级管理员密钥验证失败"}
+            
+            # 获取所有refresh token
+            with self.token_manager.token_lock:
+                refresh_tokens = list(self.token_manager.tokens.keys())
+            
+            if not refresh_tokens:
+                return {"success": False, "message": "没有可导出的token"}
+            
+            # 创建分号分割的token字符串
+            token_string = ";".join(refresh_tokens)
+            
+            # 生成建议的文件名（带时间戳）
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            suggested_filename = f"refresh_tokens_export_{timestamp}.txt"
+            
+            logger.info(f"🔒 超级管理员请求导出 {len(refresh_tokens)} 个refresh token")
+            
+            return {
+                "success": True,
+                "message": f"准备导出 {len(refresh_tokens)} 个token",
+                "content": token_string,
+                "suggested_filename": suggested_filename,
+                "token_count": len(refresh_tokens)
+            }
+            
+        except Exception as e:
+            logger.error(f"❌ 准备导出refresh token失败: {e}")
+            return {"success": False, "message": f"导出失败: {str(e)}"}
+    
+    def batch_get_refresh_tokens(self, email_url_dict: Dict[str, str], max_workers: int = 5) -> Dict[str, Any]:
+        """批量获取refresh token并自动创建用户"""
+        try:
+            from login_client import LoginClient
+            login_client = LoginClient()
+            
+            # 传递token_manager参数，这样在获取refresh_token后会立即尝试创建用户
+            results = login_client.batch_process_emails(email_url_dict, max_workers, self.token_manager)
+            
+            # 提取有效的token并添加到管理器
+            valid_tokens = []
+            for email, result in results.items():
+                if result.get('refresh_token'):
+                    valid_tokens.append(result['refresh_token'])
+            
+            if valid_tokens:
+                self.token_manager.add_refresh_tokens(valid_tokens)
+                logger.info(f"✅ 批量获取并添加了 {len(valid_tokens)} 个有效token")
+            
+            return {
+                'success': True,
+                'results': results,
+                'total_count': len(email_url_dict),
+                'success_count': len(valid_tokens)
+            }
+            
+        except Exception as e:
+            logger.error(f"❌ 批量获取refresh token失败: {e}")
+            return {'success': False, 'message': str(e)}
+    
+    def start_services(self):
+        """启动后台服务"""
+        try:
+            self.token_manager.start_auto_refresh()
+            logger.success("✅ 后台服务启动成功")
+        except Exception as e:
+            logger.error(f"❌ 启动后台服务失败: {e}")
+    
+    def stop_services(self):
+        """停止后台服务"""
+        try:
+            self.token_manager.stop_auto_refresh()
+            logger.info("⏹️ 后台服务已停止")
+        except Exception as e:
+            logger.error(f"❌ 停止后台服务失败: {e}")