mirror of
https://github.com/fastapiadmin/FastapiAdmin.git
synced 2026-09-28 06:57:54 +00:00
zhangtao
ecec6ceadb
refactor: 重构基础模型和响应类，增加success字段 fix: 修复定时任务日志记录和异常处理问题 style: 调整中间件日志记录逻辑，优化参数处理 docs: 更新数据库配置和依赖项说明 perf: 优化文件上传服务，支持OSS存储 test: 添加定时任务日志模型和参数验证 chore: 更新依赖项，添加openai库支持
37 lines
1.4 KiB
Python
37 lines
1.4 KiB
Python
from fastapi import FastAPI
|
|
from starlette.websockets import WebSocket, WebSocketDisconnect
|
|
|
|
from mcp_server.mcp_client import MCPClient
|
|
|
|
|
|
async def init_ai_websocket(app: FastAPI):
|
|
|
|
@app.websocket("/ws/chat")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
await websocket.accept()
|
|
user_id = id(websocket)
|
|
user_contexts = {}
|
|
user_contexts[user_id] = [{"role": "system", "content": "你是一个有帮助的助手。"}]
|
|
|
|
client = MCPClient()
|
|
await client.connect_to_server('mcp_server/mcp_server.py')
|
|
try:
|
|
while True:
|
|
user_msg = await websocket.receive_text()
|
|
user_contexts[user_id].append({"role": "user", "content": user_msg})
|
|
await websocket.send_json({"role": "user", "content": user_msg})
|
|
|
|
assistant_reply = ""
|
|
response = client.put_query(user_msg)
|
|
await websocket.send_json({"start": True})
|
|
async for content_piece in response:
|
|
assistant_reply += content_piece
|
|
await websocket.send_json({"role": "assistant", "content": content_piece})
|
|
|
|
await websocket.send_json({"done": True})
|
|
|
|
except WebSocketDisconnect:
|
|
print("WebSocket 断开连接")
|
|
# 清理上下文
|
|
user_contexts.pop(user_id, None)
|
|
await client.cleanup() |