Răsfoiți Sursa

增加excel merge websocket demo

zhangqian 1 an în urmă
părinte
comite
82dc74877d
1 a modificat fișierele cu 42 adăugiri și 6 ștergeri
  1. 42 6
      app/api/excel.py

+ 42 - 6
app/api/excel.py

@@ -1,14 +1,18 @@
-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()
 
 ALLOWED_EXTENSIONS = {'xlsx'}
@@ -79,7 +83,8 @@ async def run_conformity_api():
         # 运行方法
         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 @@ async def download_excel():
     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")