from fastapi import FastAPI, File, UploadFile, HTTPException, Depends from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse, StreamingResponse import uvicorn import os from typing import List import shutil from io import BytesIO import json from config import config from models import ( ChatRequest, ChatResponse, StreamChatRequest, StreamChatChunk, DocumentInfo, ErrorResponse, SuccessResponse, ) from services import AsyncRAGService from utils import ( extract_text_from_pdf_async, validate_file_size, ensure_directory_exists, is_supported_file_type, ) # 创建FastAPI应用 app = FastAPI( title=config.APP_NAME, version=config.APP_VERSION, description="高效简洁的RAG服务API", docs_url="/docs", redoc_url="/redoc", ) # 添加CORS中间件 app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 确保上传目录存在 ensure_directory_exists(config.UPLOAD_DIR) # 创建RAG服务实例 rag_service = AsyncRAGService() def get_rag_service() -> AsyncRAGService: """依赖注入:获取RAG服务实例""" return rag_service @app.get("/", response_model=dict) async def root(): """根路径 - 服务健康检查""" return { "message": f"欢迎使用 {config.APP_NAME}", "version": config.APP_VERSION, "status": "running", } @app.get("/health") async def health_check(): """健康检查接口""" return {"status": "healthy", "service": config.APP_NAME} @app.post("/upload", response_model=SuccessResponse) async def upload_document( file: UploadFile = File(...), service: AsyncRAGService = Depends(get_rag_service) ): """上传文档接口""" try: # 验证文件类型 if not is_supported_file_type(file.filename): raise HTTPException( status_code=400, detail="不支持的文件类型。目前支持:PDF, TXT" ) # 验证文件大小 content = await file.read() if not validate_file_size(len(content), config.MAX_FILE_SIZE): raise HTTPException( status_code=400, detail=f"文件过大。最大支持 {config.MAX_FILE_SIZE // 1024 // 1024}MB", ) # 提取文本内容 if file.filename.lower().endswith(".pdf"): text_content = await extract_text_from_pdf_async(BytesIO(content)) else: # txt文件 text_content = content.decode("utf-8") if not text_content.strip(): raise HTTPException(status_code=400, detail="文件内容为空或无法提取文本") # 添加到向量库 doc_id = await service.add_document_async(text_content, file.filename) # 保存文件到本地(可选) file_path = os.path.join(config.UPLOAD_DIR, f"{doc_id}_{file.filename}") with open(file_path, "wb") as f: f.write(content) return SuccessResponse( message="文档上传成功", data={ "document_id": doc_id, "filename": file.filename, "size": len(content), }, ) except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"文档处理失败: {str(e)}") @app.post("/chat", response_model=ChatResponse) async def chat( request: ChatRequest, service: AsyncRAGService = Depends(get_rag_service) ): """聊天问答接口""" try: result = await service.chat_async( question=request.question, top_k=request.top_k, temperature=request.temperature, ) return ChatResponse( answer=result["answer"], sources=result["sources"], processing_time=result["processing_time"], ) except Exception as e: raise HTTPException(status_code=500, detail=f"问答处理失败: {str(e)}") @app.post("/chat/stream") async def chat_stream( request: StreamChatRequest, service: AsyncRAGService = Depends(get_rag_service) ): """流式聊天问答接口""" async def generate_stream(): try: async for chunk_data in service.chat_stream_async( question=request.question, top_k=request.top_k, temperature=request.temperature, ): # 将数据转换为 JSON 格式并添加换行符 chunk = StreamChatChunk(**chunk_data) yield f"data: {chunk.model_dump_json()}\n\n" except Exception as e: # 发生错误时发送错误信息 error_chunk = StreamChatChunk( content=f"生成回答时发生错误: {str(e)}", is_final=True, sources=[], processing_time=0.0, ) yield f"data: {error_chunk.model_dump_json()}\n\n" return StreamingResponse( generate_stream(), media_type="text/plain", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "Access-Control-Allow-Origin": "*", "Access-Control-Allow-Headers": "*", }, ) @app.get("/documents", response_model=List[DocumentInfo]) async def get_documents(service: AsyncRAGService = Depends(get_rag_service)): """获取文档列表接口""" try: docs = await service.get_documents_async() return [ DocumentInfo( id=doc["id"], filename=doc["filename"], upload_time=doc["upload_time"], size=0, # 可以后续添加文件大小信息 chunks_count=doc["chunks_count"], ) for doc in docs ] except Exception as e: raise HTTPException(status_code=500, detail=f"获取文档列表失败: {str(e)}") @app.delete("/documents/{doc_id}", response_model=SuccessResponse) async def delete_document( doc_id: str, service: AsyncRAGService = Depends(get_rag_service) ): """删除文档接口""" try: success = await service.delete_document_async(doc_id) if not success: raise HTTPException(status_code=404, detail="文档不存在") return SuccessResponse(message="文档删除成功") except HTTPException: raise except Exception as e: raise HTTPException(status_code=500, detail=f"删除文档失败: {str(e)}") @app.exception_handler(Exception) async def global_exception_handler(request, exc): """全局异常处理器""" return JSONResponse( status_code=500, content=ErrorResponse( error="内部服务器错误", detail=str(exc) if config.DEBUG else "请联系管理员" ).dict(), ) if __name__ == "__main__": # 验证配置 try: config.validate() except ValueError as e: print(f"配置错误: {e}") exit(1) # 启动服务 uvicorn.run( "main:app", host=config.HOST, port=config.PORT, reload=config.DEBUG, log_level="info", )