|  |  |  | 
|---|
|  |  |  | import random | 
|---|
|  |  |  | import string | 
|---|
|  |  |  |  | 
|---|
|  |  |  | from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends | 
|---|
|  |  |  | from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends, Request, WebSocket | 
|---|
|  |  |  | from fastapi.responses import JSONResponse, FileResponse | 
|---|
|  |  |  | from sqlalchemy.orm import Session | 
|---|
|  |  |  | from starlette.websockets import WebSocket | 
|---|
|  |  |  | # from starlette.websockets import WebSocket | 
|---|
|  |  |  |  | 
|---|
|  |  |  | from app.api import get_current_user, get_current_user_websocket | 
|---|
|  |  |  | from app.api import get_current_user, get_current_user_websocket, Response | 
|---|
|  |  |  | from app.models import UserModel, AgentType | 
|---|
|  |  |  | from app.models.base_model import get_db | 
|---|
|  |  |  | from app.service.session import SessionService | 
|---|
|  |  |  | 
|---|
|  |  |  | return prefix + random_part | 
|---|
|  |  |  |  | 
|---|
|  |  |  |  | 
|---|
|  |  |  | def db_create_session(db: Session, user_id: str): | 
|---|
|  |  |  | def db_create_session(db: Session, user_id: str, message:str, upload_filenames: list): | 
|---|
|  |  |  | db_id = generate_db_id() | 
|---|
|  |  |  | session = SessionService(db).create_session( | 
|---|
|  |  |  | db_id, | 
|---|
|  |  |  | "合并Excel", | 
|---|
|  |  |  | message, | 
|---|
|  |  |  | "basic_excel_merge", | 
|---|
|  |  |  | AgentType.BASIC, | 
|---|
|  |  |  | int(user_id) | 
|---|
|  |  |  | int(user_id), | 
|---|
|  |  |  | {"role": "user", "content": message, "upload_filenames": upload_filenames} | 
|---|
|  |  |  | ) | 
|---|
|  |  |  | return session | 
|---|
|  |  |  |  | 
|---|
|  |  |  |  | 
|---|
|  |  |  | @router.post('/excel/upload') | 
|---|
|  |  |  | @router.post('/excel/upload', response_model=Response) | 
|---|
|  |  |  | async def upload_file(files: list[UploadFile] = File(...), current_user: UserModel = Depends(get_current_user)): | 
|---|
|  |  |  | user_id = str(current_user.id) | 
|---|
|  |  |  | if not any(file.filename for file in files): | 
|---|
|  |  |  | return JSONResponse(content={"error": "没有文件部分"}, status_code=400) | 
|---|
|  |  |  | return Response(code=400, msg="没有文件部分", data={}) | 
|---|
|  |  |  | if not user_id: | 
|---|
|  |  |  | return JSONResponse(content={"error": "缺少参数user_id"}, status_code=400) | 
|---|
|  |  |  | return Response(code=400, msg="缺少参数user_id", data={}) | 
|---|
|  |  |  | user_source = user_file_path(user_id, SOURCE_FILES_PATH) | 
|---|
|  |  |  | user_excel = 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): | 
|---|
|  |  |  | save_path = os.path.join(user_source, 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={"code": 200, "msg": "", "data": {}}, status_code=200) | 
|---|
|  |  |  | return Response(code=400, msg="不允许的文件类型", data={}) | 
|---|
|  |  |  | return Response(code=200, msg="上传成功", data={}) | 
|---|
|  |  |  |  | 
|---|
|  |  |  |  | 
|---|
|  |  |  | # ws://localhost:9201/api/document/ws/excel | 
|---|
|  |  |  | 
|---|
|  |  |  | user_excel = EXCEL_FILES_PATH | 
|---|
|  |  |  | create_dir_if_not_exists(user_source) | 
|---|
|  |  |  | create_dir_if_not_exists(user_excel) | 
|---|
|  |  |  |  | 
|---|
|  |  |  | while True: | 
|---|
|  |  |  | data = await websocket.receive_text() | 
|---|
|  |  |  | # data = await websocket.receive_text()git | 
|---|
|  |  |  | receive_message = await websocket.receive_json() | 
|---|
|  |  |  | try: | 
|---|
|  |  |  | if data == "\"合并Excel\"": | 
|---|
|  |  |  | if receive_message.get("message") == "合并Excel": | 
|---|
|  |  |  | upload_filenames = receive_message.get('upload_filenames', []) | 
|---|
|  |  |  | merge_file = run_conformity(user_source, user_excel) | 
|---|
|  |  |  | if merge_file is not None: | 
|---|
|  |  |  |  | 
|---|
|  |  |  | await websocket.send_json({ | 
|---|
|  |  |  | "message": "文档合并成功!", | 
|---|
|  |  |  | "type": "stream", | 
|---|
|  |  |  | "file_name": f"{merge_file}.xlsx", | 
|---|
|  |  |  | "download_url": f"./api/document/download/{merge_file}.xlsx" | 
|---|
|  |  |  | "files": [ | 
|---|
|  |  |  | { | 
|---|
|  |  |  | "file_name": "Excel", | 
|---|
|  |  |  | "file_url": f"./api/document/download/{merge_file}.xlsx?file_type=excel", | 
|---|
|  |  |  | } | 
|---|
|  |  |  | ] | 
|---|
|  |  |  | }) | 
|---|
|  |  |  | await websocket.send_json({ | 
|---|
|  |  |  | "message": "文档合并成功!", | 
|---|
|  |  |  | "message": "合并成功", | 
|---|
|  |  |  | "type": "close", | 
|---|
|  |  |  | }) | 
|---|
|  |  |  | # 创建会话记录 | 
|---|
|  |  |  | session = db_create_session(db, user_id) | 
|---|
|  |  |  | session = db_create_session(db, user_id, receive_message.get("message"), upload_filenames) | 
|---|
|  |  |  | # 更新会话记录 | 
|---|
|  |  |  | if session: | 
|---|
|  |  |  | session_id = session.id | 
|---|
|  |  |  | new_message = { | 
|---|
|  |  |  | "role": "user", | 
|---|
|  |  |  | "download_url": f"./api/document/download/{merge_file}.xlsx" | 
|---|
|  |  |  | "role": "assistant", | 
|---|
|  |  |  | "content": { | 
|---|
|  |  |  | "message": "\u5408\u5e76\u6210\u529f", | 
|---|
|  |  |  | "type": "message", | 
|---|
|  |  |  | "file_name": "Excel", | 
|---|
|  |  |  | "file_url": f"/api/document/download/{merge_file}.xlsx?file_type=excel" | 
|---|
|  |  |  | } | 
|---|
|  |  |  | } | 
|---|
|  |  |  | session_service = SessionService(db) | 
|---|
|  |  |  | session_service.update_session(session_id, message=new_message) | 
|---|
|  |  |  | 
|---|
|  |  |  | await websocket.send_json({"error": "合并失败", "type": "stream", "files": []}) | 
|---|
|  |  |  | await websocket.close() | 
|---|
|  |  |  | else: | 
|---|
|  |  |  | print(f"Received data: {data}") | 
|---|
|  |  |  | await websocket.send_json({"error": "未知指令", "data": str(data)}) | 
|---|
|  |  |  | print(f"Received data: {receive_message.get('message')}") | 
|---|
|  |  |  | await websocket.send_json({"error": "未知指令", "data": str(receive_message.get('message'))}) | 
|---|
|  |  |  | await websocket.close() | 
|---|
|  |  |  | except Exception as e: | 
|---|
|  |  |  | await websocket.send_json({"error": str(e)}) | 
|---|
|  |  |  | 
|---|
|  |  |  | return JSONResponse(content={"error": "文件不存在"}, status_code=404) | 
|---|
|  |  |  | return FileResponse( | 
|---|
|  |  |  | path=file_path, | 
|---|
|  |  |  | filename=file_name, | 
|---|
|  |  |  | filename="Excel.xlsx", | 
|---|
|  |  |  | media_type='application/octet-stream', | 
|---|
|  |  |  | ) | 
|---|
|  |  |  | # def delete_file(): | 
|---|