From c966f133bea46f5d1b804296b270955478fa0ac4 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期二, 18 二月 2025 13:42:35 +0800 Subject: [PATCH] 报表合并问题处理 --- app/api/excel.py | 243 ++++++++++++++++++++++++++++-------------------- 1 files changed, 140 insertions(+), 103 deletions(-) diff --git a/app/api/excel.py b/app/api/excel.py index 1ff4741..1c0ccf4 100644 --- a/app/api/excel.py +++ b/app/api/excel.py @@ -1,14 +1,15 @@ -from fastapi import APIRouter, File, UploadFile, 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 +import random +import string -from app.api import get_current_user_websocket -from app.models.agent_model import AgentModel, AgentType +from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends, Request, WebSocket +from fastapi.responses import JSONResponse, FileResponse +from sqlalchemy.orm import Session +# from starlette.websockets import WebSocket + +from app.api import get_current_user, get_current_user_websocket, Response +from app.models import UserModel, AgentType from app.models.base_model import get_db -from app.models.user_model import UserModel +from app.service.session import SessionService from app.utils.excelmerge.conformity import run_conformity import shutil import os @@ -18,124 +19,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) + + +def generate_db_id(prefix: str = "me") -> str: + random_part = ''.join(random.choices(string.ascii_letters + string.digits, k=13)) + return prefix + random_part + + +def db_create_session(db: Session, user_id: str, message:str, upload_filenames: list): + db_id = generate_db_id() + session = SessionService(db).create_session( + db_id, + message, + "basic_excel_merge", + AgentType.BASIC, + int(user_id), + {"role": "user", "content": message, "upload_filenames": upload_filenames} + ) + return session + + +@router.post('/excel/upload', response_model=Response) +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 Response(code=400, msg="娌℃湁鏂囦欢閮ㄥ垎", data={}) + if not user_id: + return Response(code=400, msg="缂哄皯鍙傛暟user_id", data={}) + user_source = user_file_path(user_id, SOURCE_FILES_PATH) + user_excel = EXCEL_FILES_PATH + + create_dir_if_not_exists(user_source) + create_dir_if_not_exists(user_excel) + clear_directory(user_source) 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 Response(code=400, msg="涓嶅厑璁哥殑鏂囦欢绫诲瀷", data={}) + return Response(code=200, msg="涓婁紶鎴愬姛", data={}) -@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), + db: Session = Depends(get_db)): 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 = EXCEL_FILES_PATH + create_dir_if_not_exists(user_source) + create_dir_if_not_exists(user_excel) + while True: + # data = await websocket.receive_text()git + receive_message = await websocket.receive_json() + try: + if receive_message.get("message") == "鍚堝苟Excel": + upload_filenames = receive_message.get('upload_filenames', []) + merge_file = run_conformity(user_source, user_excel) + if merge_file is not None: + + await websocket.send_json({ + "type": "stream", + "files": [ + { + "file_name": "Excel", + "file_url": f"./api/document/download/{merge_file}.xlsx?file_type=excel", + } + ] + }) + await websocket.send_json({ + "message": "鍚堝苟鎴愬姛", + "type": "close", + }) + # 鍒涘缓浼氳瘽璁板綍 + session = db_create_session(db, user_id, receive_message.get("message"), upload_filenames) + # 鏇存柊浼氳瘽璁板綍 + if session: + session_id = session.id + new_message = { + "role": "assistant", + "content": { + "message": "\u5408\u5e76\u6210\u529f", + "type": "message", + "file_name": "Excel", + "file_url": f"/api/document/download/{merge_file}.xlsx?file_type=excel" + } + } + session_service = SessionService(db) + session_service.update_session(session_id, message=new_message) + else: + await websocket.send_json({"error": "鍚堝苟澶辫触", "type": "stream", "files": []}) + await websocket.close() + else: + print(f"Received data: {receive_message.get('message')}") + await websocket.send_json({"error": "鏈煡鎸囦护", "data": str(receive_message.get('message'))}) + await websocket.close() + except Exception as e: + await websocket.send_json({"error": str(e)}) + await websocket.close() + + +@router.get("/download/{file_full_name}") +async def download_file(file_full_name: str): + file_name = os.path.basename(file_full_name) + user_excel = EXCEL_FILES_PATH + file_path = os.path.join(user_excel, file_full_name) + + if not os.path.exists(file_path): + return JSONResponse(content={"error": "鏂囦欢涓嶅瓨鍦�"}, status_code=404) + return FileResponse( + path=file_path, + filename="Excel.xlsx", + media_type='application/octet-stream', + ) + # def delete_file(): + # try: + # os.unlink(file_path) + # except OSError as e: + # print(f"Deleting file error") + + # 寰呬笅杞藉畬鎴愬悗鍒犻櫎鐢熸垚鐨勬枃浠� + # background_tasks.add_task(delete_file) + # return FileResponse(path=file_path, filename=file_name, + # media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") -- Gitblit v1.8.0