From 9a676f601cd6dbeba131095915ef48f1eafb12b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Wed, 9 Jul 2025 00:39:54 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=A0=BC=E5=BC=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 48 ++++++++++++++++++++++++++---------------------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/main.py b/main.py index df9ede4..e64783c 100644 --- a/main.py +++ b/main.py @@ -34,7 +34,7 @@ setup_logger( name=config.APP_NAME, level=config.LOG_LEVEL, log_dir=config.LOG_DIR, - use_async=True # 使用异步日志,避免阻塞 + use_async=True, # 使用异步日志,避免阻塞 ) logger = get_logger(__name__) @@ -77,21 +77,21 @@ async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(secur """验证Bearer token""" if not config.AUTH_ENABLED: return True - + if credentials is None: raise HTTPException( status_code=401, detail="认证失败:缺少Authorization header", - headers={"WWW-Authenticate": "Bearer"} + headers={"WWW-Authenticate": "Bearer"}, ) - + if credentials.credentials != config.API_TOKEN: raise HTTPException( status_code=401, detail="认证失败:无效的token", - headers={"WWW-Authenticate": "Bearer"} + headers={"WWW-Authenticate": "Bearer"}, ) - + return True @@ -115,16 +115,16 @@ async def health_check(): @app.post("/upload", response_model=SuccessResponse) async def upload_document( - file: UploadFile = File(...), + file: UploadFile = File(...), service: AsyncRAGService = Depends(get_rag_service), - _: bool = Depends(verify_token) + _: bool = Depends(verify_token), ): """上传文档接口""" start_time = time.time() - + try: logger.info(f"开始上传文档: {file.filename}") - + # 验证文件类型 if not is_supported_file_type(file.filename): logger.warning(f"不支持的文件类型: {file.filename}") @@ -160,7 +160,9 @@ async def upload_document( f.write(content) duration = time.time() - start_time - logger.info(f"文档上传成功: {file.filename}, 文档ID: {doc_id}, 耗时: {duration:.2f}s") + logger.info( + f"文档上传成功: {file.filename}, 文档ID: {doc_id}, 耗时: {duration:.2f}s" + ) return SuccessResponse( message="文档上传成功", @@ -175,22 +177,25 @@ async def upload_document( raise except Exception as e: duration = time.time() - start_time - logger.error(f"文档上传失败: {file.filename}, 错误: {str(e)}, 耗时: {duration:.2f}s", exc_info=True) + logger.error( + f"文档上传失败: {file.filename}, 错误: {str(e)}, 耗时: {duration:.2f}s", + exc_info=True, + ) raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}") @app.post("/chat", response_model=ChatResponse) async def chat( - request: ChatRequest, + request: ChatRequest, service: AsyncRAGService = Depends(get_rag_service), - _: bool = Depends(verify_token) + _: bool = Depends(verify_token), ): """聊天问答接口""" start_time = time.time() - + try: logger.info(f"开始处理问答: {request.question[:50]}...") - + result = await service.chat_async( question=request.question, top_k=request.top_k, @@ -214,9 +219,9 @@ async def chat( @app.post("/chat/stream") async def chat_stream( - request: StreamChatRequest, + request: StreamChatRequest, service: AsyncRAGService = Depends(get_rag_service), - _: bool = Depends(verify_token) + _: bool = Depends(verify_token), ): """流式聊天问答接口""" logger.info(f"开始处理流式问答: {request.question[:50]}...") @@ -257,8 +262,7 @@ async def chat_stream( @app.get("/documents", response_model=List[DocumentInfo]) async def get_documents( - service: AsyncRAGService = Depends(get_rag_service), - _: bool = Depends(verify_token) + service: AsyncRAGService = Depends(get_rag_service), _: bool = Depends(verify_token) ): """获取文档列表接口""" try: @@ -282,9 +286,9 @@ async def get_documents( @app.delete("/documents/{doc_id}", response_model=SuccessResponse) async def delete_document( - doc_id: str, + doc_id: str, service: AsyncRAGService = Depends(get_rag_service), - _: bool = Depends(verify_token) + _: bool = Depends(verify_token), ): """删除文档接口""" try: