Просмотр исходного кода

改为websocket 脚本运行。(功能有待优化)

zhangxiao 1 год назад
Родитель
Сommit
a77073fec8
1 измененных файлов с 50 добавлено и 83 удалено
  1. 50 83
      app/api/excel.py

+ 50 - 83
app/api/excel.py

@@ -4,7 +4,6 @@ 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
@@ -18,8 +17,6 @@ router = APIRouter()
 ALLOWED_EXTENSIONS = {'xlsx'}
 EXCEL_FILES_PATH = 'data/output'
 SOURCE_FILES_PATH = 'data/source'
-output_path_value = None
-
 
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
@@ -30,23 +27,29 @@ def create_dir_if_not_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):
+    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:
@@ -63,79 +66,43 @@ async def upload_file(files: list[UploadFile] = File(...)):
     return JSONResponse(content={"message": "文件上传成功", "paths": save_path_list}, status_code=201)
 
 
-@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)
+# ws://localhost:9201/api/document/ws/excel
+@router.websocket("/ws/excel")
+async def ws_excel(websocket: WebSocket):
+    await websocket.accept()
+    while True:
+        data = await websocket.receive_json()
+        action = data.get("action")
+        try:
+            if action == "process":
+                clear_directory(EXCEL_FILES_PATH)
+                output_file_path = run_conformity()
+                await websocket.send_json({"step_message": "开始合并"})
+            elif action == "inquire":
+                files = os.listdir(EXCEL_FILES_PATH)
+                if not files:
+                    await websocket.send_json({"step_message": "正在合并中"})
+                else:
+                    await websocket.send_json({"step_message": "文档合并成功!"})
+            elif action == "download":
+                files = os.listdir(EXCEL_FILES_PATH)
+                if not files:
+                    await websocket.send_json({"error": "目录下没有生成的文件"})
+                else:
+                    first_file = files[0]
+                    await websocket.send_json({"step_message": "合并文件已生成", "download_url": f"/download/{first_file}"})
+            else:
+                await websocket.send_json({"error": "未知指令"})
+        except Exception as e:
+            await websocket.send_json({"error": str(e)})
 
 
-@router.get('/excel/download_excel')
-async def download_excel():
+@router.get("/download/{filename}")
+async def download_file(filename: str):
     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")
+        raise HTTPException(status_code=500, detail="服务器错误")