mirror of
https://github.com/fastapiadmin/FastapiAdmin.git
synced 2026-09-23 21:18:09 +00:00
zhangtao
f57d4f8cfb
本次提交包含多项优化和清理工作： 1. 新增订单类型枚举BUY，调整定时任务状态显示逻辑 2. 优化Redis连接配置，重构认证逻辑减少重复代码 3. 移除大量初始化演示数据、过期的租户与角色配置 4. 简化调度器状态存储，不再写入配置表而是内存缓存 5. 修复插件开关状态存储逻辑，调整字典数据标签与值 6. 优化WebSocket聊天处理，完善错误处理流程 7. 清理冗余导入与废弃的操作日志过滤逻辑
97 lines
3.5 KiB
Python
97 lines
3.5 KiB
Python
import json
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
from app.core.database import async_db_session
|
|
from app.core.dependencies import _authenticate
|
|
from app.core.exceptions import CustomException
|
|
from app.core.logger import logger
|
|
from app.core.router_class import OperationLogRoute
|
|
|
|
from .schema import ChatQuerySchema
|
|
from .service import ChatService
|
|
|
|
WS_AI = APIRouter(
|
|
route_class=OperationLogRoute,
|
|
prefix="/ai/chat",
|
|
tags=["智能助手WebSocket"],
|
|
)
|
|
|
|
|
|
async def _send_error_and_close(websocket: WebSocket, message: str) -> None:
|
|
"""发送错误消息并关闭连接"""
|
|
try:
|
|
await websocket.send_text(f"错误: {message}")
|
|
except RuntimeError:
|
|
pass
|
|
finally:
|
|
try:
|
|
await websocket.close()
|
|
except RuntimeError:
|
|
pass
|
|
|
|
|
|
@WS_AI.websocket("/ws", name="WebSocket聊天")
|
|
async def websocket_chat_controller(websocket: WebSocket) -> None:
|
|
"""
|
|
WebSocket 聊天接口。
|
|
|
|
支持两种消息格式：
|
|
1. 纯文本：直接发送消息内容
|
|
2. JSON 格式：{"message": "消息内容", "session_id": "会话ID", "files": [...]}
|
|
|
|
ws://127.0.0.1:8001/api/v1/ai/chat/ws?token=xxx
|
|
"""
|
|
await websocket.accept()
|
|
token = websocket.query_params.get("token")
|
|
|
|
if not token:
|
|
await _send_error_and_close(websocket, "未提供认证token，请重新登录")
|
|
return
|
|
|
|
try:
|
|
# 认证：db 会话需在整个连接生命周期内保持打开（auth.db 供 ChatService 使用）
|
|
redis = websocket.app.state.redis
|
|
async with async_db_session() as db:
|
|
auth = await _authenticate(token, db, redis)
|
|
|
|
user = auth.user
|
|
logger.info("WebSocket连接已建立: {} - 用户: {}", websocket.client, user.username if user else "未认证")
|
|
|
|
# 消息循环
|
|
while True:
|
|
try:
|
|
data = await websocket.receive_text()
|
|
try:
|
|
message_data = json.loads(data)
|
|
query = ChatQuerySchema(**message_data)
|
|
logger.info("收到聊天查询: {} - 会话ID: {}", query, query.session_id)
|
|
|
|
chat_result = ChatService.chat_query(query=query, auth=auth)
|
|
async for chunk in chat_result:
|
|
if chunk:
|
|
try:
|
|
await websocket.send_text(chunk)
|
|
except RuntimeError:
|
|
logger.warning("WebSocket连接已关闭，停止发送消息")
|
|
return
|
|
except json.JSONDecodeError:
|
|
logger.warning("收到非JSON消息: {}", data)
|
|
await websocket.send_text("消息格式错误，请发送JSON格式的消息")
|
|
except Exception as e:
|
|
logger.error("处理消息时出错: {}", e)
|
|
await websocket.send_text(f"处理消息时出错: {e}")
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info("WebSocket连接已断开: {}", websocket.client)
|
|
return
|
|
|
|
except CustomException as e:
|
|
# 认证失败等业务异常
|
|
logger.warning("WebSocket认证失败: {}", e.msg)
|
|
await _send_error_and_close(websocket, e.msg)
|
|
except Exception as e:
|
|
# 未知异常
|
|
logger.exception("WebSocket未知异常: {}", e)
|
|
await _send_error_and_close(websocket, "服务器内部错误")
|