增加excel merge websocket demo

zhangqian

2024-10-16 82dc74877d0757ad4562ead7bc68e45f65a8ff14

增加excel merge websocket demo

 app/api/excel.py

@@ -1,13 +1,17 @@
from fastapi import APIRouter, File, UploadFile
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

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.utils.excelmerge.conformity import run_conformity
from pathlib import Path
import subprocess
import shutil
import os


router = APIRouter()

@@ -79,7 +83,8 @@
        # 运行方法
        output_path = run_conformity()
        output_path_value = output_path
        return JSONResponse(content={"message": "conformity.py 运行成功", "output_path": str(output_path)}, status_code=200)
        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)

@@ -102,4 +107,35 @@
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="文件不存在")
    except Exception as e:
        raise HTTPException(status_code=500, detail="服务器错误")
        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")