Merge branch 'master' of http://192.168.5.5:10010/r/rag-gateway

zhaoqingang

2024-12-18 916d984d2628fd8d863183cf5ba9e5a0a7339871

 app/api/excel.py

@@ -1,13 +1,7 @@
from fastapi import APIRouter, File, UploadFile, Depends
from fastapi import APIRouter, File, UploadFile
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 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 starlette.websockets import WebSocket
from app.utils.excelmerge.conformity import run_conformity
import shutil
import os
@@ -17,6 +11,7 @@
ALLOWED_EXTENSIONS = {'xlsx'}
EXCEL_FILES_PATH = 'data/output'
SOURCE_FILES_PATH = 'data/source'


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@@ -56,34 +51,54 @@
        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(SOURCE_FILES_PATH, 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)


# ws://localhost:9201/api/document/ws/excel
@router.websocket("/ws/excel")
async def ws_excel(websocket: WebSocket):
    await websocket.accept()

    create_dir_if_not_exists(SOURCE_FILES_PATH)
    create_dir_if_not_exists(EXCEL_FILES_PATH)

    while True:
        data = await websocket.receive_text()
        try:
            if data == "合并Excel文件":
                clear_directory(EXCEL_FILES_PATH)
                output_file_path = run_conformity()
                await websocket.send_json({"step_message": "开始合并", "type": "stream", "files": []})
            elif data == "查询合并进度":
            if data == "\"合并Excel\"":
                run_excel = run_conformity()
                files = os.listdir(EXCEL_FILES_PATH)
                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(EXCEL_FILES_PATH)
                if not files:
                    await websocket.send_json({"step_message": "正在合并中", "type": "stream", "files": []})
                else:
                    await websocket.send_json({"step_message": "文档合并成功！", "type": "stream", "files": []})
            elif data == "获取文件":
            elif data == "\"获取文件\"":
                files = os.listdir(EXCEL_FILES_PATH)
                if not files:
                    await websocket.send_json({"error": "目录下没有生成的文件", "type": "stream", "files": []})
@@ -100,9 +115,11 @@
                        }]
                    })
            else:
                await websocket.send_json({"error": "未知指令"})
                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/{filename}")
@@ -113,4 +130,4 @@
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="文件不存在")
    except Exception as e:
        raise HTTPException(status_code=500, detail="服务器错误")
        raise HTTPException(status_code=500, detail="服务器错误")