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

zhaoqingang

2024-12-18 916d984d2628fd8d863183cf5ba9e5a0a7339871

 app/api/excel.py

@@ -1,14 +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
@@ -18,7 +11,6 @@
ALLOWED_EXTENSIONS = {'xlsx'}
EXCEL_FILES_PATH = 'data/output'
SOURCE_FILES_PATH = 'data/source'
output_path_value = None


def allowed_file(filename):
@@ -30,112 +22,112 @@
        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):
    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": "目录已清空"}


@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)
    create_dir_if_not_exists(EXCEL_FILES_PATH)
    clear_directory(SOURCE_FILES_PATH)
    clear_directory(EXCEL_FILES_PATH)

    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(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)


@router.post('/excel/conformity')
async def run_conformity_api():
    global output_path_value  # 声明全局变量
# 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\"":
                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 == "\"获取文件\"":
                files = os.listdir(EXCEL_FILES_PATH)
                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/{filename}")
async def download_file(filename: str):
    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,
        return FileResponse(os.path.join(EXCEL_FILES_PATH, filename), filename=filename,
                            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)

    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:

        print(f"Client {chat_id} disconnected")