Files
FastapiAdmin/backend/app/plugin/module_application/ai/ws.py
T
zhangtao 1d54dca76c style: 统一代码格式和字符串引号使用
refactor: 优化代码结构和可读性

feat: 添加http_limit模块实现请求限制功能

fix: 修复异步任务中使用time.sleep的问题

chore: 更新依赖项并添加pytest测试框架

docs: 更新项目描述信息

perf: 优化Redis序列化方式使用JSON替代pickle

test: 添加测试相关配置和依赖
2026-01-17 20:07:25 +08:00

46 lines
1.5 KiB
Python

from fastapi import APIRouter, WebSocket
from app.core.logger import log
from app.core.router_class import OperationLogRoute
from .schema import ChatQuerySchema
from .service import McpService
WS_AI = APIRouter(
route_class=OperationLogRoute,
prefix="/application/ai",
tags=["MCP智能助手WebSocket"],
)
@WS_AI.websocket("/ws", name="WebSocket聊天")
async def websocket_chat_controller(
websocket: WebSocket,
) -> None:
"""
WebSocket聊天接口
ws://127.0.0.1:8001/api/v1/application/ai/ws
"""
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
# 流式发送响应
try:
async for chunk in McpService.chat_query(query=ChatQuerySchema(message=data)):
if chunk:
await websocket.send_text(chunk)
except Exception as e:
log.error(f"处理聊天查询出错: {e!s}")
await websocket.send_text(f"抱歉,处理您的请求时出现了错误: {e!s}")
except Exception as e:
log.error(f"WebSocket聊天出错: {e!s}")
finally:
try:
# 检查WebSocket连接状态,避免重复关闭已关闭的连接
if websocket.client_state != websocket.client_state.DISCONNECTED:
await websocket.close()
except Exception as e:
log.debug(f"WebSocket关闭时发生异常(预期行为,服务可能正在关闭): {e!s}")