100 lines
4.1 KiB
Python
100 lines
4.1 KiB
Python
"""WebSocket 路由：流式对话 + 重新生成。协议：
|
|
客户端 → {"type":"chat","content":"...","agent_id":null,"model":"","file_ids":[]}
|
|
{"type":"regenerate"}
|
|
服务端 → {"type":"delta","message_id":1,"content":"增量文本"}
|
|
{"type":"done","message_id":1}
|
|
{"type":"suggestions","message_id":1,"items":["..",".."]}
|
|
{"type":"title","title":"..."}
|
|
{"type":"error","message":"..."}
|
|
"""
|
|
import asyncio
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from sqlalchemy.orm import Session
|
|
|
|
from ..core.security import decode_token
|
|
from ..database import SessionLocal
|
|
from ..models import ChatSession, User
|
|
from ..services import chat_service
|
|
|
|
router = APIRouter()
|
|
|
|
# session_id -> set[WebSocket]（用于推送建议/标题）
|
|
_conns: dict[int, set[WebSocket]] = {}
|
|
|
|
|
|
def push_to_session(session_id: int, payload: dict):
|
|
"""向该会话的所有在线连接推送消息。"""
|
|
for ws in list(_conns.get(session_id, set())):
|
|
try:
|
|
asyncio.get_event_loop().create_task(ws.send_json(payload))
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@router.websocket("/chat/ws")
|
|
async def chat_ws(websocket: WebSocket):
|
|
await websocket.accept()
|
|
token = websocket.query_params.get("token", "")
|
|
session_id = websocket.query_params.get("session_id", "")
|
|
payload = decode_token(token)
|
|
if not payload:
|
|
await websocket.send_json({"type": "error", "message": "认证失败"})
|
|
await websocket.close()
|
|
return
|
|
|
|
db: Session = SessionLocal()
|
|
try:
|
|
user = db.get(User, int(payload["sub"]))
|
|
if not user or not user.is_active:
|
|
await websocket.send_json({"type": "error", "message": "用户无效"})
|
|
await websocket.close()
|
|
return
|
|
|
|
session = db.get(ChatSession, int(session_id)) if session_id.isdigit() else None
|
|
if not session or session.user_id != user.id:
|
|
await websocket.send_json({"type": "error", "message": "会话不存在"})
|
|
await websocket.close()
|
|
return
|
|
|
|
_conns.setdefault(session.id, set()).add(websocket)
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
msg_type = data.get("type")
|
|
|
|
if msg_type == "chat":
|
|
content = (data.get("content") or "").strip()
|
|
file_ids = data.get("file_ids") or []
|
|
if not content and not file_ids:
|
|
await websocket.send_json({"type": "error", "message": "消息不能为空"})
|
|
continue
|
|
try:
|
|
async for delta, message_id in chat_service.chat_stream(
|
|
db, user, session, content,
|
|
agent_id=data.get("agent_id"), model=data.get("model", ""),
|
|
file_ids=file_ids,
|
|
):
|
|
await websocket.send_json({"type": "delta", "message_id": message_id, "content": delta})
|
|
await websocket.send_json({"type": "done", "message_id": message_id})
|
|
except Exception as e:
|
|
await websocket.send_json({"type": "error", "message": f"模型调用失败：{e}"})
|
|
|
|
elif msg_type == "regenerate":
|
|
try:
|
|
async for delta, message_id in chat_service.regenerate_stream(db, session):
|
|
await websocket.send_json({"type": "delta", "message_id": message_id, "content": delta})
|
|
await websocket.send_json({"type": "done", "message_id": message_id})
|
|
except LookupError as e:
|
|
await websocket.send_json({"type": "error", "message": str(e)})
|
|
except Exception as e:
|
|
await websocket.send_json({"type": "error", "message": f"模型调用失败：{e}"})
|
|
finally:
|
|
_conns.get(session.id, set()).discard(websocket)
|
|
if not _conns.get(session.id):
|
|
_conns.pop(session.id, None)
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
db.close()
|