""" 霜云(Shimokumo) - API服务器模块 基于Flask的REST API服务器,提供以下端点: - POST /chat : 聊天对话(支持流式) - POST /search : 联网搜索 - POST /browse : 网页浏览 - POST /novel : 小说创作管理 - POST /video : 视频制作管理 - POST /media : 多媒体播放管理 - GET /health : 健康检查 - GET /info : 模型信息 支持流式响应(SSE)、API Key认证和CORS。 """ import json import time from functools import wraps from typing import Any, Dict, Optional from flask import Flask, Response, jsonify, request, stream_with_context from configuration_shimokumo import ShimokumoConfig from shimokumo_utils.logger import get_logger logger = get_logger("Shimokumo.Server") def create_api_key_auth(api_key: str): """创建API Key认证装饰器工厂""" def decorator(f): @wraps(f) def decorated_function(*args, **kwargs): if not api_key: return f(*args, **kwargs) # 从header或query参数获取API Key provided_key = request.headers.get("X-API-Key") or request.args.get("api_key") if not provided_key: return jsonify({ "error": "未提供API Key", "message": "请在请求头中添加 X-API-Key 或在URL参数中添加 api_key", }), 401 if provided_key != api_key: return jsonify({ "error": "API Key无效", "message": "请检查您的API Key是否正确", }), 403 return f(*args, **kwargs) return decorated_function return decorator class ShimokumoServer: """霜云API服务器 封装Flask应用,提供REST API接口。 用法: server = ShimokumoServer(config, inference, ...) server.run() 或使用工厂方法: app = create_app(config, inference) app.run(host="0.0.0.0", port=8080) """ def __init__( self, config: ShimokumoConfig, inference_engine=None, web_search_module=None, web_browser_module=None, media_player_module=None, novel_creator_module=None, video_generator_module=None, character_chat_module=None, ): """ 初始化API服务器。 Args: config: 模型配置 inference_engine: 推理引擎 web_search_module: 联网搜索模块 web_browser_module: 网页浏览模块 media_player_module: 多媒体播放模块 novel_creator_module: 小说创作模块 video_generator_module: 视频制作模块 character_chat_module: 角色聊天模块 """ self.config = config self.inference = inference_engine self.web_search = web_search_module self.web_browser = web_browser_module self.media_player = media_player_module self.novel_creator = novel_creator_module self.video_generator = video_generator_module self.character_chat = character_chat_module self.app = self._create_flask_app() self.start_time = time.time() def _create_flask_app(self) -> Flask: """创建并配置Flask应用""" app = Flask( __name__, static_folder=None, template_folder=None, ) # CORS支持 self._setup_cors(app) # API Key认证装饰器 auth = create_api_key_auth(self.config.api_key) # ==================== 健康检查 ==================== @app.route("/health", methods=["GET"]) def health_check(): """健康检查端点""" uptime = time.time() - self.start_time return jsonify({ "status": "healthy", "service": "Shimokumo", "uptime": uptime, "version": "1.0.0", }) @app.route("/info", methods=["GET"]) def model_info(): """模型信息端点""" info = { "name": self.config.character_name, "name_en": self.config.character_name_en, "title": self.config.character_title, "model_config": { "hidden_size": self.config.hidden_size, "num_layers": self.config.num_layers, "num_heads": self.config.num_heads, "num_kv_heads": self.config.num_kv_heads, "vocab_size": self.config.vocab_size, "max_seq_len": self.config.max_seq_len, }, "capabilities": { "web_search": self.config.enable_web_search, "novel_creation": self.config.enable_novel_creation, "video_generation": self.config.enable_video_generation, "character_chat": self.config.enable_character_chat, "media_player": self.config.enable_media_player, "web_browser": self.config.enable_web_browser, }, } if self.inference: model = self.inference.model info["model_params"] = f"{model.get_num_params() / 1e9:.2f}B" info["model_size_mb"] = f"{model.get_model_size_mb():.1f}MB" return jsonify(info) # ==================== 聊天接口 ==================== @app.route("/chat", methods=["POST"]) @auth def chat(): """聊天对话接口 请求体: message (str): 用户消息 stream (bool): 是否流式输出 user_id (str): 用户ID(多轮对话) max_tokens (int): 最大生成token数 temperature (float): 采样温度 top_p (float): Top-P采样 """ data = request.get_json(silent=True) or {} message = data.get("message", "") is_stream = data.get("stream", False) user_id = data.get("user_id", "default") max_tokens = data.get("max_tokens", self.config.max_new_tokens) temperature = data.get("temperature", self.config.temperature) top_p = data.get("top_p", self.config.top_p) if not message: return jsonify({"error": "message不能为空"}), 400 # 角色聊天模式 if data.get("roleplay", False) and self.character_chat: response = self.character_chat.chat( message, user_id=user_id, max_new_tokens=max_tokens, ) return jsonify(response.to_dict()) # 普通聊天 if not self.inference: return jsonify({ "response": f"({self.config.character_name}还没有加载模型的说...请稍后再试)", "error": "模型未加载", }), 503 if is_stream: return Response( stream_with_context(self._stream_chat( message, user_id, max_tokens, temperature, top_p )), mimetype="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }, ) else: response = self.inference.chat_reply( user_message=message, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, ) return jsonify({ "response": response, "user_id": user_id, "model": self.config.character_name, }) # ==================== 搜索接口 ==================== @app.route("/search", methods=["POST"]) @auth def search(): """联网搜索接口 请求体: query (str): 搜索关键词 engines (list): 搜索引擎列表 max_results (int): 最大结果数 """ if not self.config.enable_web_search: return jsonify({"error": "搜索功能未启用"}), 403 data = request.get_json(silent=True) or {} query = data.get("query", "") engines = data.get("engines") max_results = data.get("max_results", 10) if not query: return jsonify({"error": "query不能为空"}), 400 results = self.web_search.search( query, engines=engines, max_total_results=max_results, ) return jsonify({ "query": query, "results": [r.to_dict() for r in results], "total": len(results), }) # ==================== 浏览接口 ==================== @app.route("/browse", methods=["POST"]) @auth def browse(): """网页浏览接口 请求体: url (str): 网页URL extract_summary (bool): 是否提取摘要 """ if not self.config.enable_web_browser: return jsonify({"error": "浏览功能未启用"}), 403 data = request.get_json(silent=True) or {} url = data.get("url", "") extract_summary = data.get("extract_summary", True) if not url: return jsonify({"error": "url不能为空"}), 400 page = self.web_browser.fetch(url) if not page: return jsonify({"error": f"无法获取网页: {url}"}), 404 result = page.to_dict() if extract_summary: result["summary"] = self.web_browser.generate_summary(page) return jsonify(result) # ==================== 小说接口 ==================== @app.route("/novel", methods=["POST"]) @auth def novel(): """小说创作接口 请求体: action (str): 操作类型 (create_project/set_world/add_character/ generate_outline/generate_chapter/check/save/load) ... (各操作的具体参数) """ if not self.config.enable_novel_creation: return jsonify({"error": "小说功能未启用"}), 403 data = request.get_json(silent=True) or {} action = data.get("action", "") if not self.novel_creator: return jsonify({"error": "小说模块未初始化"}), 503 try: result = self._handle_novel_action(data) return jsonify(result) except Exception as e: logger.error(f"小说操作失败: {e}") return jsonify({"error": str(e)}), 500 # ==================== 视频接口 ==================== @app.route("/video", methods=["POST"]) @auth def video(): """视频制作接口 请求体: action (str): 操作类型 (create/generate_storyboard/generate_subtitles/export_srt) ... (各操作的具体参数) """ if not self.config.enable_video_generation: return jsonify({"error": "视频功能未启用"}), 403 data = request.get_json(silent=True) or {} action = data.get("action", "") if not self.video_generator: return jsonify({"error": "视频模块未初始化"}), 503 try: result = self._handle_video_action(data) return jsonify(result) except Exception as e: logger.error(f"视频操作失败: {e}") return jsonify({"error": str(e)}), 500 # ==================== 媒体接口 ==================== @app.route("/media", methods=["POST"]) @auth def media(): """多媒体播放接口 请求体: action (str): 操作类型 (parse/add_playlist/get_playlist/resolve_stream) url (str): 媒体URL """ if not self.config.enable_media_player: return jsonify({"error": "媒体功能未启用"}), 403 data = request.get_json(silent=True) or {} action = data.get("action", "") url = data.get("url", "") if not self.media_player: return jsonify({"error": "媒体模块未初始化"}), 503 if action == "parse" and url: item = self.media_player.parse_url(url) return jsonify(item.to_dict()) elif action == "add_playlist" and url: item = self.media_player.add_to_playlist(url, data.get("title", "")) return jsonify(item.to_dict()) elif action == "get_playlist": return jsonify({ "name": self.media_player.playlist.name, "items": len(self.media_player.playlist), "info": self.media_player.get_playlist_info(), }) elif action == "resolve_stream" and url: stream_url = self.media_player.resolve_stream_url(url) return jsonify({"stream_url": stream_url}) else: return jsonify({"error": "未知的操作类型或缺少参数"}), 400 # ==================== 错误处理 ==================== @app.errorhandler(404) def not_found(e): return jsonify({ "error": "未找到请求的资源", "message": f"请检查您的请求路径。可用端点: /chat, /search, /browse, /novel, /video, /media, /health, /info", }), 404 @app.errorhandler(500) def server_error(e): return jsonify({ "error": "服务器内部错误", "message": str(e), }), 500 return app def _setup_cors(self, app: Flask) -> None: """配置CORS跨域支持""" @app.after_request def add_cors_headers(response): response.headers.add("Access-Control-Allow-Origin", "*") response.headers.add("Access-Control-Allow-Headers", "Content-Type, X-API-Key") response.headers.add("Access-Control-Allow-Methods", "GET, POST, OPTIONS") return response # 处理OPTIONS预检请求 @app.before_request def handle_options(): if request.method == "OPTIONS": response = app.make_default_options_response() response.headers.add("Access-Control-Allow-Origin", "*") response.headers.add("Access-Control-Allow-Headers", "Content-Type, X-API-Key") response.headers.add("Access-Control-Allow-Methods", "GET, POST, OPTIONS") return response def _stream_chat( self, message: str, user_id: str, max_tokens: int, temperature: float, top_p: float, ): """SSE流式聊天生成器""" try: for token_text in self.inference.chat_stream( user_message=message, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, ): # SSE格式: data: {json}\n\n chunk = json.dumps({"token": token_text}, ensure_ascii=False) yield f"data: {chunk}\n\n" # 发送结束标记 yield "data: [DONE]\n\n" except Exception as e: error_chunk = json.dumps({"error": str(e)}, ensure_ascii=False) yield f"data: {error_chunk}\n\n" def _handle_novel_action(self, data: Dict[str, Any]) -> Dict[str, Any]: """处理小说创作操作""" action = data.get("action") creator = self.novel_creator if action == "create_project": project = creator.create_project( title=data.get("title", "未命名"), genre=data.get("genre", "奇幻"), ) return {"success": True, "project_id": project.id, "title": project.title} elif action == "set_world": project_id = data.get("project_id") # 查找项目 project = None for p in creator.projects.values(): if p.id == project_id: project = p break if not project: return {"error": "项目不存在"} creator.set_world(project, **data) return {"success": True} elif action == "generate_outline": return {"error": "需要提供project_id"} elif action == "save": return {"success": True} else: return {"error": f"未知的操作: {action}"} def _handle_video_action(self, data: Dict[str, Any]) -> Dict[str, Any]: """处理视频制作操作""" action = data.get("action") generator = self.video_generator if action == "create": project = generator.create_project( title=data.get("title", "未命名视频"), width=data.get("width", 1920), height=data.get("height", 1080), fps=data.get("fps", 24), ) return {"success": True, "project_id": project.id} elif action == "generate_subtitles": return {"error": "需要提供project_id"} elif action == "export_srt": return {"error": "需要提供project_id"} else: return {"error": f"未知的操作: {action}"} def run(self, host: Optional[str] = None, port: Optional[int] = None, debug: bool = False): """ 启动服务器。 Args: host: 监听地址(覆盖配置) port: 监听端口(覆盖配置) debug: 调试模式 """ host = host or self.config.host port = port or self.config.port logger.info(f"霜云API服务器启动中...") logger.info(f"地址: http://{host}:{port}") logger.info(f"API文档: http://{host}:{port}/info") logger.info(f"健康检查: http://{host}:{port}/health") self.app.run( host=host, port=port, debug=debug or self.config.debug, threaded=True, )