""" API处理模块 处理主要的API路由逻辑 """ import json import time import logging from typing import Dict, List from fastapi import HTTPException, Request from fastapi.responses import StreamingResponse, JSONResponse from src.config import Config from src.constants import ( APIConstants, ResponseConstants, LogMessages, ErrorMessages, HeaderConstants ) from src.exceptions import ( AuthenticationError, SerializationError, K2ThinkProxyError ) from src.models import ChatCompletionRequest, ModelsResponse, ModelInfo from src.tool_handler import ToolHandler from src.response_processor import ResponseProcessor logger = logging.getLogger(__name__) class APIHandler: """API处理器""" def __init__(self, config: Config): self.config = config self.tool_handler = ToolHandler(config) self.response_processor = ResponseProcessor(config, self.tool_handler) def validate_api_key(self, authorization: str) -> bool: """验证API密钥""" if not authorization or not authorization.startswith(APIConstants.BEARER_PREFIX): return False api_key = authorization[APIConstants.BEARER_PREFIX_LENGTH:] # 移除 "Bearer " 前缀 return api_key == self.config.VALID_API_KEY def should_output_thinking(self, model_name: str) -> bool: """根据模型名判断是否应该输出思考内容""" return model_name != APIConstants.MODEL_ID_NOTHINK def get_actual_model_id(self, model_name: str) -> str: """获取实际的模型ID（将nothink版本映射回原始模型）""" if model_name == APIConstants.MODEL_ID_NOTHINK: return APIConstants.MODEL_ID return model_name async def get_models(self) -> ModelsResponse: """获取模型列表""" model_info_standard = ModelInfo( id=APIConstants.MODEL_ID, created=int(time.time()), owned_by=APIConstants.MODEL_OWNER, root=APIConstants.MODEL_ROOT ) model_info_nothink = ModelInfo( id=APIConstants.MODEL_ID_NOTHINK, created=int(time.time()), owned_by=APIConstants.MODEL_OWNER, root=APIConstants.MODEL_ROOT ) return ModelsResponse(data=[model_info_standard, model_info_nothink]) async def chat_completions(self, request: ChatCompletionRequest, auth_request: Request): """处理聊天补全请求""" # 验证API密钥 authorization = auth_request.headers.get(HeaderConstants.AUTHORIZATION, "") if not self.validate_api_key(authorization): raise AuthenticationError() # 判断是否应该输出思考内容 output_thinking = self.should_output_thinking(request.model) actual_model_id = self.get_actual_model_id(request.model) try: # 处理消息 raw_messages = self._process_raw_messages(request.messages) # 检查工具是否启用和存在 has_tools = self._check_tools_enabled(request) self._log_request_info(raw_messages, has_tools, request.tools) # 处理工具相关消息 processed_messages = self._process_messages_with_tools( raw_messages, request, has_tools ) # 构建K2Think请求 k2think_payload = self._build_k2think_payload( request, processed_messages, actual_model_id ) # 验证JSON序列化 self._validate_json_serialization(k2think_payload) # 设置请求头 headers = self._build_request_headers(request, k2think_payload) # 处理响应 if request.stream: return await self._handle_stream_response( k2think_payload, headers, has_tools, output_thinking, request.model ) else: return await self._handle_non_stream_response( k2think_payload, headers, has_tools, output_thinking, request.model ) except K2ThinkProxyError: # 重新抛出自定义异常 raise except Exception as e: logger.error(f"API转发错误: {e}") raise HTTPException( status_code=APIConstants.HTTP_INTERNAL_ERROR, detail={ "error": { "message": str(e), "type": ErrorMessages.API_ERROR } } ) def _process_raw_messages(self, messages: List) -> List[Dict]: """处理原始消息""" raw_messages = [] for msg in messages: try: raw_messages.append({ "role": msg.role, "content": msg.content, # 保持原始格式，稍后再转换 "tool_calls": msg.tool_calls }) except Exception as e: logger.error(f"处理消息时出错: {e}, 消息: {msg}") # 使用默认值 raw_messages.append({ "role": msg.role, "content": str(msg.content) if msg.content else "", "tool_calls": msg.tool_calls }) return raw_messages def _check_tools_enabled(self, request: ChatCompletionRequest) -> bool: """检查工具是否启用""" return ( self.config.TOOL_SUPPORT and request.tools is not None and len(request.tools) > 0 and request.tool_choice != "none" ) def _log_request_info(self, raw_messages: List[Dict], has_tools: bool, tools: List): """记录请求信息""" logger.info(LogMessages.TOOL_STATUS.format( has_tools, len(tools) if tools else 0 )) logger.info(LogMessages.MESSAGE_RECEIVED.format(len(raw_messages))) # 记录原始消息的角色分布 role_count = {} for msg in raw_messages: role = msg.get("role", "unknown") role_count[role] = role_count.get(role, 0) + 1 logger.info(LogMessages.ROLE_DISTRIBUTION.format("原始", role_count)) def _process_messages_with_tools( self, raw_messages: List[Dict], request: ChatCompletionRequest, has_tools: bool ) -> List[Dict]: """处理工具相关消息""" if has_tools: processed_messages = self.tool_handler.process_messages_with_tools( raw_messages, request.tools, request.tool_choice ) logger.info(LogMessages.MESSAGE_PROCESSED.format( len(raw_messages), len(processed_messages) )) # 记录处理后消息的角色分布 processed_role_count = {} for msg in processed_messages: role = msg.get("role", "unknown") processed_role_count[role] = processed_role_count.get(role, 0) + 1 logger.info(LogMessages.ROLE_DISTRIBUTION.format("处理后", processed_role_count)) else: processed_messages = raw_messages logger.info(LogMessages.NO_TOOLS) return processed_messages def _build_k2think_payload( self, request: ChatCompletionRequest, processed_messages: List[Dict], actual_model_id: str = None ) -> Dict: """构建K2Think请求负载""" # 构建K2Think格式的请求体 - 支持多模态内容 k2think_messages = [] for msg in processed_messages: try: # 使用多模态内容转换函数 content = self.response_processor.content_to_multimodal(msg.get("content", "")) k2think_messages.append({ "role": msg["role"], "content": content }) except Exception as e: logger.error(f"构建K2Think消息时出错: {e}, 消息: {msg}") # 使用安全的默认值 fallback_content = self.tool_handler._content_to_string(msg.get("content", "")) k2think_messages.append({ "role": msg.get("role", "user"), "content": fallback_content }) # 使用实际的模型ID model_id = actual_model_id or APIConstants.MODEL_ID return { "stream": request.stream, "model": model_id, "messages": k2think_messages, "params": {}, "tool_servers": [], "features": { "image_generation": False, "code_interpreter": False, "web_search": False }, "variables": self.response_processor.get_current_datetime_info(), "model_item": { "id": model_id, "object": ResponseConstants.MODEL_OBJECT, "owned_by": APIConstants.MODEL_OWNER, "root": APIConstants.MODEL_ROOT, "parent": None, "status": "active", "connection_type": "external", "name": model_id }, "background_tasks": { "title_generation": True, "tags_generation": True }, "chat_id": self.response_processor.generate_chat_id(), "id": self.response_processor.generate_session_id(), "session_id": self.response_processor.generate_session_id() } def _validate_json_serialization(self, k2think_payload: Dict): """验证JSON序列化""" try: # 测试JSON序列化 json.dumps(k2think_payload, ensure_ascii=False) logger.info(LogMessages.JSON_VALIDATION_SUCCESS) except Exception as e: logger.error(LogMessages.JSON_VALIDATION_FAILED.format(e)) # 尝试修复序列化问题 try: k2think_payload = json.loads(json.dumps(k2think_payload, default=str, ensure_ascii=False)) logger.info(LogMessages.JSON_FIXED) except Exception as fix_error: logger.error(f"无法修复序列化问题: {fix_error}") raise SerializationError() def _build_request_headers(self, request: ChatCompletionRequest, k2think_payload: Dict) -> Dict[str, str]: """构建请求头""" return { HeaderConstants.ACCEPT: ( HeaderConstants.EVENT_STREAM_JSON if request.stream else HeaderConstants.APPLICATION_JSON ), HeaderConstants.CONTENT_TYPE: HeaderConstants.APPLICATION_JSON, HeaderConstants.AUTHORIZATION: f"{APIConstants.BEARER_PREFIX}{self.config.K2THINK_TOKEN}", HeaderConstants.ORIGIN: "https://www.k2think.ai", HeaderConstants.REFERER: "https://www.k2think.ai/c/" + k2think_payload["chat_id"], HeaderConstants.USER_AGENT: HeaderConstants.DEFAULT_USER_AGENT } async def _handle_stream_response( self, k2think_payload: Dict, headers: Dict[str, str], has_tools: bool, output_thinking: bool = True, original_model: str = None ) -> StreamingResponse: """处理流式响应""" return StreamingResponse( self.response_processor.process_stream_response_with_tools( k2think_payload, headers, has_tools, output_thinking, original_model ), media_type=HeaderConstants.TEXT_EVENT_STREAM, headers={ HeaderConstants.CACHE_CONTROL: HeaderConstants.NO_CACHE, HeaderConstants.CONNECTION: HeaderConstants.KEEP_ALIVE, HeaderConstants.X_ACCEL_BUFFERING: HeaderConstants.NO_BUFFERING } ) async def _handle_non_stream_response( self, k2think_payload: Dict, headers: Dict[str, str], has_tools: bool, output_thinking: bool = True, original_model: str = None ) -> JSONResponse: """处理非流式响应""" full_content, token_info = await self.response_processor.process_non_stream_response( k2think_payload, headers, output_thinking ) # 处理工具调用 tool_calls = None message_content = full_content if has_tools: tool_calls = self.tool_handler.extract_tool_invocations(full_content) if tool_calls: # 当存在工具调用时，内容必须为null（OpenAI规范） message_content = None logger.info(LogMessages.TOOL_CALLS_EXTRACTED.format( json.dumps(tool_calls, ensure_ascii=False) )) else: # 从内容中移除工具JSON message_content = self.tool_handler.remove_tool_json_content(full_content) if not message_content: message_content = full_content # 保留原内容如果清理后为空 openai_response = self.response_processor.create_completion_response( message_content, tool_calls, token_info, original_model ) return JSONResponse(content=openai_response)