From 9191081e44363e4227ea3230edd6fb7619ccb294 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期四, 26 十二月 2024 13:51:32 +0800 Subject: [PATCH] Merge branch 'master' of http://192.168.5.5:10010/r/rag-gateway --- app/api/excel.py | 231 ++++++++++++++++++++++++++++++++------------------------- 1 files changed, 131 insertions(+), 100 deletions(-) diff --git a/app/api/excel.py b/app/api/excel.py index 1ff4741..9c13103 100644 --- a/app/api/excel.py +++ b/app/api/excel.py @@ -1,14 +1,9 @@ -from fastapi import APIRouter, File, UploadFile, Depends +from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends from fastapi.responses import JSONResponse, FileResponse -from fastapi.exceptions import HTTPException -from sqlalchemy.orm import Session -from starlette.websockets import WebSocket, WebSocketDisconnect -from werkzeug.utils import secure_filename +from starlette.websockets import WebSocket -from app.api import get_current_user_websocket -from app.models.agent_model import AgentModel, AgentType -from app.models.base_model import get_db -from app.models.user_model import UserModel +from app.api import get_current_user, get_current_user_websocket +from app.models import UserModel from app.utils.excelmerge.conformity import run_conformity import shutil import os @@ -18,124 +13,160 @@ ALLOWED_EXTENSIONS = {'xlsx'} EXCEL_FILES_PATH = 'data/output' SOURCE_FILES_PATH = 'data/source' -output_path_value = None -def allowed_file(filename): +def allowed_file(filename: str) -> bool: return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS -def create_dir_if_not_exists(path): +def create_dir_if_not_exists(path: str): if not os.path.exists(path): os.makedirs(path) -@router.post('/excel/upload') -async def upload_file(files: list[UploadFile] = File(...)): - if not any(file.filename for file in files): - return JSONResponse(content={"error": "娌℃湁鏂囦欢閮ㄥ垎"}, status_code=400) - - create_dir_if_not_exists(SOURCE_FILES_PATH) - - # 娓呯┖SOURCE_FILES_PATH鐩綍 - for filename in os.listdir(SOURCE_FILES_PATH): - file_path = os.path.join(SOURCE_FILES_PATH, filename) +def clear_directory(path: str) -> dict: + for filename in os.listdir(path): + file_path = os.path.join(path, filename) try: if os.path.isfile(file_path) or os.path.islink(file_path): os.unlink(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path) except Exception as e: - return JSONResponse(content={"error": "鏂囦欢澶勭悊鍑洪敊"}, status_code=500) + return {"error": "娓呯┖鍑洪敊"} + return {"message": "鐩綍宸叉竻绌�"} + + +def user_file_path(userid: str, path: str) -> str: + return os.path.join(path, userid) + + +@router.post('/excel/upload') +async def upload_file(files: list[UploadFile] = File(...), current_user: UserModel = Depends(get_current_user)): + user_id = str(current_user.id) + if not any(file.filename for file in files): + return JSONResponse(content={"error": "娌℃湁鏂囦欢閮ㄥ垎"}, status_code=400) + if not user_id: + return JSONResponse(content={"error": "缂哄皯鍙傛暟user_id"}, status_code=400) + user_source = user_file_path(user_id, SOURCE_FILES_PATH) + user_excel = user_file_path(user_id, EXCEL_FILES_PATH) + + create_dir_if_not_exists(user_source) + create_dir_if_not_exists(user_excel) + clear_directory(user_source) + clear_directory(user_excel) save_path_list = [] for file in files: if file.filename == '': return JSONResponse(content={"error": "娌℃湁閫夋嫨鏂囦欢"}, status_code=400) if file and allowed_file(file.filename): - filename = secure_filename(file.filename) - save_path = os.path.join(SOURCE_FILES_PATH, filename) + save_path = os.path.join(user_source, file.filename) with open(save_path, 'wb') as buffer: shutil.copyfileobj(file.file, buffer) save_path_list.append(save_path) else: return JSONResponse(content={"error": "涓嶅厑璁哥殑鏂囦欢绫诲瀷"}, status_code=400) - return JSONResponse(content={"message": "鏂囦欢涓婁紶鎴愬姛", "paths": save_path_list}, status_code=201) + return JSONResponse(content={"code": 200, "msg": "", "data": {}}, status_code=200) -@router.post('/excel/conformity') -async def run_conformity_api(): - global output_path_value # 澹版槑鍏ㄥ眬鍙橀噺 - try: - create_dir_if_not_exists(EXCEL_FILES_PATH) - - # 娓呯┖EXCEL_FILES_PATH鐩綍 - for filename in os.listdir(EXCEL_FILES_PATH): - file_path = os.path.join(EXCEL_FILES_PATH, filename) - try: - if os.path.isfile(file_path) or os.path.islink(file_path): - os.unlink(file_path) - elif os.path.isdir(file_path): - shutil.rmtree(file_path) - except Exception as e: - return JSONResponse(content={"error": "鏂囦欢澶勭悊鍑洪敊"}, status_code=500) - - # 杩愯鏂规硶 - output_path = run_conformity() - output_path_value = output_path - return JSONResponse(content={"message": "conformity.py 杩愯鎴愬姛", "output_path": str(output_path)}, - status_code=200) - except Exception as e: - return JSONResponse(content={"error": str(e)}, status_code=500) - - -@router.get('/excel/file/status') -async def get_file_status(): - try: - return JSONResponse(content={"output_path": str(output_path_value)}, status_code=200) - except Exception as e: - return JSONResponse(content={"error": str(e)}, status_code=500) - - -@router.get('/excel/download_excel') -async def download_excel(): - try: - files = os.listdir(EXCEL_FILES_PATH) - first_file = files[0] - return FileResponse(os.path.join(EXCEL_FILES_PATH, first_file), filename=first_file, - media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet') - except FileNotFoundError: - raise HTTPException(status_code=404, detail="鏂囦欢涓嶅瓨鍦�") - except Exception as e: - raise HTTPException(status_code=500, detail="鏈嶅姟鍣ㄩ敊璇�") - - -@router.websocket("/ws/{agent_id}/{chat_id}") -async def excel_chat(websocket: WebSocket, - agent_id: str, - chat_id: str, - db: Session = Depends(get_db)): - agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first() - if not agent: - ret = {"message": "Agent not found", "type": "close"} - return websocket.send_json(ret) - agent_type = agent.agent_type - if chat_id == "" or chat_id == "0": - ret = {"message": "Chat ID not found", "type": "close"} - return websocket.send_json(ret) - - if agent_type != AgentType.BASIC: - ret = {"message": "agent type error", "type": "close"} - return websocket.send_json(ret) - +# ws://localhost:9201/api/document/ws/excel +@router.websocket("/ws/excel") +async def ws_excel(websocket: WebSocket, current_user: UserModel = Depends(get_current_user_websocket)): await websocket.accept() - try: - while True: - message = await websocket.receive_json() - print(message) # 鎵撳嵃鎺ユ敹鍒扮殑娑堟伅 - result = {"message": "宸茬敓鎴愭枃浠�", "type": "file", "url": "ip/download?id=xxxx"} - # 鍙戦�佸搷搴� - await websocket.send_json(result) - except WebSocketDisconnect as e: + user_id = str(current_user.id) - print(f"Client {chat_id} disconnected") + user_source = user_file_path(user_id, SOURCE_FILES_PATH) + user_excel = user_file_path(user_id, EXCEL_FILES_PATH) + create_dir_if_not_exists(user_source) + create_dir_if_not_exists(user_excel) + + while True: + data = await websocket.receive_text() + try: + if data == "\"鍚堝苟Excel\"": + run_excel = run_conformity(user_source, user_excel) + files = os.listdir(user_excel) + if run_excel: + first_file = files[0] + file_name = os.path.basename(first_file) + download_url = f"./api/document/download/{first_file}" + await websocket.send_json({ + "message": "鏂囨。鍚堝苟鎴愬姛锛�", + "type": "stream", + "files": [{ + "file_name": file_name, + "file_url": download_url + }] + }) + await websocket.send_json({ + "message": "鏂囨。鍚堝苟鎴愬姛锛�", + "type": "close", + }) + else: + await websocket.send_json({"error": "鍚堝苟澶辫触", "type": "stream", "files": []}) + elif data == "\"鏌ヨ鍚堝苟杩涘害\"": + files = os.listdir(user_excel) + if not files: + await websocket.send_json({"step_message": "姝e湪鍚堝苟涓�", "type": "stream", "files": []}) + else: + await websocket.send_json({"step_message": "鏂囨。鍚堝苟鎴愬姛锛�", "type": "stream", "files": []}) + elif data == "\"鑾峰彇鏂囦欢\"": + files = os.listdir(user_excel) + if not files: + await websocket.send_json({"error": "鐩綍涓嬫病鏈夌敓鎴愮殑鏂囦欢", "type": "stream", "files": []}) + else: + first_file = files[0] + file_name = os.path.basename(first_file) + file_url = f"./api/document/download/{first_file}" + await websocket.send_json({ + "step_message": "鏂囨。鍚堝苟鎴愬姛锛�", + "type": "stream", + "files": [{ + "file_name": file_name, + "file_url": file_url + }] + }) + else: + print(f"Received data: {data}") + await websocket.send_json({"error": "鏈煡鎸囦护", "data": str(data)}) + except Exception as e: + await websocket.send_json({"error": str(e)}) + await websocket.close() + + +@router.get("/download/excel") +async def download_file(background_tasks: BackgroundTasks, current_user: UserModel = Depends(get_current_user)): + user_id = str(current_user.id) + user_excel = user_file_path(user_id, EXCEL_FILES_PATH) + user_source = user_file_path(user_id, SOURCE_FILES_PATH) + + if not os.path.exists(user_excel): + return JSONResponse(status_code=404, content={"error": "鐢ㄦ埛鐩綍涓嶅瓨鍦�"}) + + excel_files = [f for f in os.listdir(user_excel) if os.path.isfile(os.path.join(user_excel, f))] + excel_files.sort(key=lambda x: os.path.getmtime(os.path.join(user_excel, x)), reverse=True) + if not excel_files: + return JSONResponse(status_code=404, content={"error": "鐢ㄦ埛鐩綍鍐呮病鏈夋枃浠�"}) + + filename = excel_files[0] + file_path = os.path.join(user_excel, filename) + + def delete_files_in_directory(directory): + for root, dirs, files in os.walk(directory, topdown=False): + for name in files: + os.remove(os.path.join(root, name)) + for name in dirs: + os.rmdir(os.path.join(root, name)) + + def delete_file(): + try: + delete_files_in_directory(user_excel) + delete_files_in_directory(user_source) + except OSError as e: + print(f"Error deleting file {file_path}: {e}") + + background_tasks.add_task(delete_file) + + return FileResponse(file_path, filename=filename, + media_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet') -- Gitblit v1.8.0