| | |
| | | import random |
| | | import string |
| | | |
| | | from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends |
| | | from fastapi.responses import JSONResponse, FileResponse |
| | | from sqlalchemy.orm import Session |
| | | from starlette.websockets import WebSocket |
| | | |
| | | from app.api import get_current_user, get_current_user_websocket |
| | | from app.models import UserModel |
| | | from app.models import UserModel, AgentType |
| | | from app.models.base_model import get_db |
| | | from app.service.session import SessionService |
| | | from app.utils.excelmerge.conformity import run_conformity |
| | | import shutil |
| | | import os |
| | |
| | | return os.path.join(path, userid) |
| | | |
| | | |
| | | def generate_db_id(prefix: str = "me") -> str: |
| | | random_part = ''.join(random.choices(string.ascii_letters + string.digits, k=13)) |
| | | return prefix + random_part |
| | | |
| | | |
| | | def db_create_session(db: Session, user_id: str): |
| | | db_id = generate_db_id() |
| | | session = SessionService(db).create_session( |
| | | db_id, |
| | | "合并Excel", |
| | | "basic_excel_merge", |
| | | AgentType.BASIC, |
| | | int(user_id) |
| | | ) |
| | | return session |
| | | |
| | | |
| | | @router.post('/excel/upload') |
| | | async def upload_file(files: list[UploadFile] = File(...), current_user: UserModel = Depends(get_current_user)): |
| | | user_id = str(current_user.id) |
| | |
| | | |
| | | # ws://localhost:9201/api/document/ws/excel |
| | | @router.websocket("/ws/excel") |
| | | async def ws_excel(websocket: WebSocket, current_user: UserModel = Depends(get_current_user_websocket)): |
| | | async def ws_excel(websocket: WebSocket, |
| | | current_user: UserModel = Depends(get_current_user_websocket), |
| | | db: Session = Depends(get_db)): |
| | | await websocket.accept() |
| | | user_id = str(current_user.id) |
| | | |
| | |
| | | "message": "文档合并成功!", |
| | | "type": "close", |
| | | }) |
| | | # 创建会话记录 |
| | | session = db_create_session(db, user_id) |
| | | # 更新会话记录 |
| | | if session: |
| | | session_id = session.id |
| | | new_message = { |
| | | "role": "user", |
| | | "download_url": f"./api/document/download/{merge_file}.xlsx" |
| | | } |
| | | session_service = SessionService(db) |
| | | session_service.update_session(session_id, message=new_message) |
| | | else: |
| | | 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)}) |
| | | await websocket.close() |
| | | except Exception as e: |
| | | await websocket.send_json({"error": str(e)}) |
| | | await websocket.close() |
| | |
| | | if not os.path.exists(file_path): |
| | | return JSONResponse(content={"error": "文件不存在"}, status_code=404) |
| | | |
| | | def delete_file(): |
| | | try: |
| | | os.unlink(file_path) |
| | | except OSError as e: |
| | | print(f"Deleting file error") |
| | | # def delete_file(): |
| | | # try: |
| | | # os.unlink(file_path) |
| | | # except OSError as e: |
| | | # print(f"Deleting file error") |
| | | |
| | | # 待下载完成后删除生成的文件 |
| | | background_tasks.add_task(delete_file) |
| | | return FileResponse(path=file_path, filename=file_name, |
| | | media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") |
| | | # background_tasks.add_task(delete_file) |
| | | # return FileResponse(path=file_path, filename=file_name, |
| | | # media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") |