rag-gateway.git

			@@ -1,13 +1,15 @@
			from fastapi import APIRouter, File, UploadFile, Depends
			import random
			import string

			from fastapi import APIRouter, File, UploadFile, Form, BackgroundTasks, Depends, Request
			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 starlette.websockets import 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.models.user_model import UserModel
			from app.service.session import SessionService
			from app.utils.excelmerge.conformity import run_conformity
			import shutil
			import os
			@@ -18,17 +20,17 @@
			EXCEL_FILES_PATH = 'data/output'
			SOURCE_FILES_PATH = 'data/source'

			def allowed_file(filename):

			def allowed_file(filename: str) -> bool:
			return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


			def create_dir_if_not_exists(path):
			def create_dir_if_not_exists(path: str):
			if not os.path.exists(path):
			os.makedirs(path)


			# 清理函数
			def clear_directory(path):
			def clear_directory(path: str) -> dict:
			for filename in os.listdir(path):
			file_path = os.path.join(path, filename)
			try:
			@@ -41,76 +43,124 @@
			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)
			def user_file_path(userid: str, path: str) -> str:
			return os.path.join(path, userid)

			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)

			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', 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 Response(code=400, msg="没有文件部分", data={})
			if not user_id:
			return Response(code=400, msg="缺少参数user_id", data={})
			user_source = user_file_path(user_id, SOURCE_FILES_PATH)
			user_excel = EXCEL_FILES_PATH

			create_dir_if_not_exists(user_source)
			create_dir_if_not_exists(user_excel)
			clear_directory(user_source)

			save_path_list = []
			for file in files:
			if file.filename == '':
			return JSONResponse(content={"error": "没有选择文件"}, status_code=400)
			if file and allowed_file(file.filename):
			filename = secure_filename(file.filename)
			save_path = os.path.join(SOURCE_FILES_PATH, 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={"message": "文件上传成功", "paths": save_path_list}, status_code=201)
			return Response(code=400, msg="不允许的文件类型", data={})
			return Response(code=200, msg="上传成功", data={})


			# ws://localhost:9201/api/document/ws/excel
			@router.websocket("/ws/excel")
			async def ws_excel(websocket: 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)

			user_source = user_file_path(user_id, SOURCE_FILES_PATH)
			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()
			try:
			if data == "合并Excel文件":
			clear_directory(EXCEL_FILES_PATH)
			output_file_path = run_conformity()
			await websocket.send_json({"step_message": "开始合并", "type": "stream", "files": []})
			elif data == "查询合并进度":
			files = os.listdir(EXCEL_FILES_PATH)
			if not files:
			await websocket.send_json({"step_message": "正在合并中", "type": "stream", "files": []})
			else:
			await websocket.send_json({"step_message": "文档合并成功！", "type": "stream", "files": []})
			elif data == "获取文件":
			files = os.listdir(EXCEL_FILES_PATH)
			if not files:
			await websocket.send_json({"error": "目录下没有生成的文件", "type": "stream", "files": []})
			else:
			first_file = files[0]
			file_name = os.path.basename(first_file)
			file_url = f"./api/document/download/{first_file}"
			if data == "\"合并Excel\"":
			merge_file = run_conformity(user_source, user_excel)
			if merge_file is not None:

			await websocket.send_json({
			"step_message": "文档合并成功！",
			"type": "stream",
			"files": [{
			"file_name": file_name,
			"file_url": file_url
			}]
			"files": [
			{
			"file_name": "Excel",
			"file_url": f"./api/document/download/{merge_file}.xlsx?file_type=excel",
			}
			]
			})
			await websocket.send_json({
			"message": "合并成功",
			"type": "close",
			})
			# 创建会话记录
			session = db_create_session(db, user_id)
			# 更新会话记录
			if session:
			session_id = session.id
			new_message = {
			"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)
			else:
			await websocket.send_json({"error": "合并失败", "type": "stream", "files": []})
			await websocket.close()
			else:
			await websocket.send_json({"error": "未知指令"})
			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()


			@router.get("/download/{filename}")
			async def download_file(filename: str):
			try:
			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.get("/download/{file_full_name}")
			async def download_file(file_full_name: str):
			file_name = os.path.basename(file_full_name)
			user_excel = EXCEL_FILES_PATH
			file_path = os.path.join(user_excel, file_full_name)

			if not os.path.exists(file_path):
			return JSONResponse(content={"error": "文件不存在"}, status_code=404)
			return FileResponse(
			path=file_path,
			filename="Excel.xlsx",
			media_type='application/octet-stream',
			)