基于WebSocket的大模型对话功能

下面是基于WebSocket实现的大模型对话功能。

安装包：

pip install fastapi 'uvicorn[standard]' dashscope shortuuid -i https://pypi.tuna.tsinghua.edu.cn/simple

代码实现：

# llm.py
import logging
import traceback

from dashscope.aigc.generation import AioGeneration


logger = logging.getLogger(__name__)


class DashScopeLLMClient:

    def __init__(self, base_url, api_key):
        self.base_url: str = base_url
        self.api_key: str = api_key
        self.default_model = 'qwen3-max'

    async def chat_stream(self, model, stream, messages):
        try:
            # 使用异步的方式调用大模型，不会阻塞接口
            response = await AioGeneration.call(
                api_key=self.api_key,
                model=model,
                stream=stream,
                messages=messages,
                result_format="message",
                incremental_output=stream
            )
            return response
        except Exception as e:
            logger.error(traceback.format_exc())

# formater.py
from typing import List, Dict

from manager import CompletionMessage


class LLMMessageFormater:

    def __init__(self, history_msg_limit=20):
        self.history_msg_limit = history_msg_limit

    async def format(self, session) -> List:
        """大模型消息格式化
        """
        llm_messages = []
        messages = session.messages
        if messages:
            messages = messages[-self.history_msg_limit:]
            for message in messages:
                llm_messages.append({
                    "role": message.role,
                    "content": message.content,
                })

        return llm_messages


class RespMessageFormater:

    @staticmethod
    async def format(session_id: str, message: CompletionMessage) -> Dict:
        return {
            "session_id": session_id,
            "message_id": message.id,
            "message": message.content,
            "timestamp": message.timestamp,
        }

# manager.py
import time
import asyncio
import datetime
import logging
import traceback
from enum import Enum
from typing import List, Dict

import shortuuid
from fastapi import WebSocket


logger = logging.getLogger(__name__)


class SessionStatus(str, Enum):
    CREATED = 'created'
    RUNNING = 'running'
    ERROR = 'error'
    CANCELLED = 'cancelled'
    COMPLETED = 'completed'


class Role(str, Enum):
    SYSTEM = 'system'
    ASSISTANT = 'assistant'
    USER = 'user'


class BaseMessage:
    pass


class CompletionMessage(BaseMessage):

    def __init__(self,
                 role: Role,
                 content: str | None,
                 name: str | None = None):
        self.id = shortuuid.uuid()
        self.name = name
        self.role = role
        self.content = content
        self.timestamp = time.time()


class Session:

    def __init__(self,
                 session_id: str | None = None,
                 task: asyncio.Task | None = None,
                 status: SessionStatus = SessionStatus.CREATED):
        self.id = session_id or shortuuid.uuid()
        self.task: asyncio.Task = task
        self.cancel_event = asyncio.Event()
        self.messages: List[CompletionMessage] = []
        self.status: SessionStatus = status
        self.is_cancelled = False
        self.start_time = datetime.datetime.now().isoformat()
        self.end_time = None


class SessionManager:

    def __init__(self):
        self.sessions: Dict[str, Session] = {}

    async def create_session(self, session_id):
        """创建会话
        """
        if not session_id:
            session = Session()
            self.sessions[session.id] = session
        elif not self.sessions.get(session_id):
            session = Session(session_id)
            self.sessions[session.id] = session
        else:
            session = self.sessions.get(session_id)

        logger.info(f"会话 {session_id} 创建成功")
        return session

    async def get_session(self, session_id: str) -> Session:
        """获取会话信息
        """
        return self.sessions.get(session_id)

    async def cancel_session(self, session_id=None):
        """取消会话
        """
        session_list = []
        if session_id:
            session = self.sessions.get(session_id)
            if not session:
                raise ValueError(f"会话 {session_id} 不存在，取消失败")

            session_list.append(session)
        else:
            if self.sessions:
                session_list = self.sessions.values()

        if session_list:
            enable_cancel_status = [SessionStatus.CREATED, SessionStatus.RUNNING]
            for session in session_list:
                if session.status in enable_cancel_status:
                    session.cancel_event.set()
                    session.is_cancelled = True
                    session.status = SessionStatus.CANCELLED
                    session.end_time = datetime.datetime.now().isoformat()

            return True
        return False


class WebSocketManager:
    def __init__(self, client_id=None):
        self.client_id = client_id or shortuuid.uuid()
        self.connections: List[WebSocket] = []
        self.session_manager: Dict[str, SessionManager] = {}

    async def connect(self, websocket: WebSocket, client_id: str):
        try:
            await websocket.accept()
        except Exception as e:
            logger.error(traceback.format_exc())
            raise

        self.connections.append(websocket)

        if not self.client_id or self.client_id not in self.session_manager:
            self.client_id = client_id
            self.session_manager[client_id] = SessionManager()

        logger.info(f"客户端 {self.client_id} 已连接")

    async def disconnect(self, websocket: WebSocket):
        # 移除WebSocket连接
        self.connections.remove(websocket)

        # 清空sessions
        session_manager = self.session_manager.pop(self.client_id)
        if session_manager:
            session_manager.sessions = []

        logger.info(f"客户端 {self.client_id} 连接已断开")

    @staticmethod
    async def send_text_message(message: str, websocket: WebSocket):
        try:
            await websocket.send_text(message)
        except Exception as e:
            logger.error(traceback.format_exc())
            raise

    @staticmethod
    async def send_json_message(message: dict, websocket: WebSocket):
        try:
            await websocket.send_json(message)
        except Exception as e:
            logger.error(traceback.format_exc())
            raise

    def check_session(self, session_id):
        if self.client_id not in self.session_manager:
            logger.warning(f"客户端 {self.client_id} 连接不存在")
            raise ValueError(f"客户端 {self.client_id} 连接不存在")

        session_manager = self.session_manager[self.client_id]
        if session_id not in session_manager.sessions:
            logger.warning(f"客户端 {self.client_id} 会话 {session_id} 不存在")
            raise ValueError(f"客户端 {self.client_id} 会话 {session_id} 不存在")

        return session_manager.sessions[session_id]

    def add_history(self, session_id: str, message: CompletionMessage):
        """添加历史消息
        """
        try:
            session = self.check_session(session_id)
            session.messages.append(message)
            return
        except Exception as e:
            logger.error(traceback.format_exc())
            raise

    def get_history(self, session_id: str) -> List[CompletionMessage]:
        """获取历史消息
        """
        try:
            session = self.check_session(session_id)
            return session.messages
        except Exception as e:
            logger.error(traceback.format_exc())
            raise

# app.py
import json
import logging
import asyncio
import traceback
from http import HTTPStatus

from fastapi.responses import HTMLResponse
from fastapi import FastAPI, WebSocket, WebSocketDisconnect

from llm import DashScopeLLMClient
from manager import WebSocketManager
from manager import Role, CompletionMessage, SessionStatus
from formater import LLMMessageFormater, RespMessageFormater


logger = logging.getLogger(__name__)


BASE_URL = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
API_KEY = 'xxx'


app = FastAPI()


websocket_manager = WebSocketManager()
llm_client = DashScopeLLMClient(base_url=BASE_URL, api_key=API_KEY)


async def completion(websocket, session, client_args):
    llm_msg_formater = LLMMessageFormater()
    resp_msg_formater = RespMessageFormater()

    session_id = session.id

    user_message = client_args.get("message", "")
    stream = client_args.get("stream", True)

    # 保存为历史消息
    message = CompletionMessage(
        name="user",
        role=Role.USER,
        content=user_message
    )
    websocket_manager.add_history(session_id, message)

    # 格式化llm消息
    llm_messages = await llm_msg_formater.format(session)
    completion_args = {
        "model": client_args.get("model", llm_client.default_model),
        "messages": llm_messages,
        "stream": stream
    }

    try:
        full_content = ""
        if stream:
            async for resp in await llm_client.chat_stream(**completion_args):
                if session.cancel_event.is_set():
                    return

                if resp.status_code == HTTPStatus.OK and resp.output.choices:
                    choices = resp.output.choices
                    if not choices:
                        continue

                    curr_content = choices[0].message.content
                    full_content += curr_content

                    resp_message = await resp_msg_formater.format(
                        session_id,
                        CompletionMessage(
                            role=Role.ASSISTANT,
                            content=curr_content
                        )
                    )

                    await websocket_manager.send_json_message({
                        **resp_message,
                        "status": SessionStatus.RUNNING
                    }, websocket)

                    if session.cancel_event.is_set():
                        return
                else:
                    resp_message = await resp_msg_formater.format(
                        session_id,
                        CompletionMessage(
                            role=Role.ASSISTANT,
                            content=resp.message
                        )
                    )
                    await websocket_manager.send_json_message({
                        **resp_message,
                        "status": SessionStatus.ERROR
                    }, websocket)

        else:
            resp = await llm_client.chat_stream(**completion_args)
            if resp.status_code == HTTPStatus.OK and resp.output.choices:
                full_content = resp.output.choices[0].message.content
                resp_message = await resp_msg_formater.format(
                    session_id,
                    CompletionMessage(
                        role=Role.ASSISTANT,
                        content=full_content
                    )
                )
                await websocket_manager.send_json_message({
                    **resp_message,
                    "status": SessionStatus.RUNNING
                }, websocket)

        if not session.is_cancelled:
            resp_message = await resp_msg_formater.format(
                session_id,
                CompletionMessage(
                    role=Role.ASSISTANT,
                    content=None
                )
            )
            await websocket_manager.send_json_message({
                **resp_message,
                "status": SessionStatus.COMPLETED
            }, websocket)

            message = CompletionMessage(
                name="assistant",
                role=Role.ASSISTANT,
                content=full_content
            )
            websocket_manager.add_history(session_id, message)
    except asyncio.CancelledError:
        raise
    except json.JSONDecodeError:
        resp_message = await resp_msg_formater.format(
            session_id,
            CompletionMessage(
                role=Role.SYSTEM,
                content="无效的JSON格式"
            )
        )

        await websocket_manager.send_json_message({
            **resp_message,
            "status": SessionStatus.ERROR
        }, websocket)
    except Exception as e:
        logger.error(traceback.format_exc())

        resp_message = await resp_msg_formater.format(
            session_id,
            CompletionMessage(
                role=Role.SYSTEM,
                content=f"处理消息时出错: {str(e)}"
            )
        )
        await websocket_manager.send_json_message({
            **resp_message,
            "status": SessionStatus.ERROR
        }, websocket)

        raise

@app.websocket("/ws/chat/{client_id}")
async def websocket_chat(websocket: WebSocket, client_id: str):
    try:
        # WebSoket连接，创建session_manager，一个client_id对应一个session_manager
        await websocket_manager.connect(websocket, client_id)
    except Exception as e:
        logger.error(traceback.format_exc())
        return

    session_manager = websocket_manager.session_manager.get(client_id)
    if not session_manager:
        return

    resp_msg_formater = RespMessageFormater()

    try:
        while True:
            # 客户端消息
            data = await websocket.receive_text()
            args = json.loads(data)
            # 通过session manager创建session
            session_id = args.get('session_id')
            session = await session_manager.create_session(session_id)
            session_id = session.id
            args['session_id'] = session_id

            user_message = args.get("message", "")

            if user_message.strip() == 'cancel':
                if not session.is_cancelled:
                    await session_manager.cancel_session(session_id)

                    resp_message = await resp_msg_formater.format(
                        session_id,
                        CompletionMessage(
                            role=Role.ASSISTANT,
                            content="会话已取消"
                        )
                    )
                    await websocket_manager.send_json_message({
                        **resp_message,
                        "status": SessionStatus.CANCELLED
                    }, websocket)
            else:
                # 如果session状态是cancelled，更新为created
                if session.is_cancelled:
                    session.status = SessionStatus.CREATED
                    session.cancel_event.clear()
                    session.is_cancelled = False

            # 创建task，并在后台运行
            task = asyncio.create_task(completion(websocket, session, args))
            session.task = task

    except asyncio.CancelledError:
        await websocket_manager.disconnect(websocket)
    except WebSocketDisconnect:
        logger.info(f"客户端 {client_id} 连接已断开")
        await websocket_manager.disconnect(websocket)
    except Exception as e:
        logger.error(traceback.format_exc())
        await websocket_manager.disconnect(websocket)


@app.get("/")
async def get():
    """聊天页面
    """
    return HTMLResponse("""
    <!DOCTYPE html>
    <html>
    <head>
        <title>聊天吧</title>
        <style>
            body { font-family: Arial, sans-serif; max-width: 800px; margin: 50px auto; }
            #messages { height: 500px; overflow-y: auto; border: 1px solid #ccc; padding: 10px; margin-bottom: 20px; }
            .message { margin: 10px 0; padding: 5px; }
            .user { background-color: #e3f2fd; text-align: right; }
            .assistant { background-color: #f5f5f5; }
            .status { color: #666; font-style: italic; }
            #messageForm { display: flex; gap: 10px; align-items: center; flex-wrap: wrap; }
            #messageInput { flex-grow: 1; padding: 10px; min-width: 200px; }
            #modelSelect { padding: 8px; border: 1px solid #ccc; border-radius: 4px; }
            .form-group { display: flex; align-items: center; gap: 5px; }
            button { padding: 10px 20px; }
            .controls { display: flex; gap: 15px; align-items: center; flex-wrap: wrap; }
            .session-info { 
                background-color: #f0f8ff; 
                padding: 8px 12px; 
                border-radius: 4px; 
                margin-bottom: 10px;
                font-size: 14px;
                display: flex;
                justify-content: space-between;
                align-items: center;
            }
            .session-id { 
                font-family: monospace; 
                color: #2c5282;
                font-weight: bold;
            }
            #clearSessionBtn {
                padding: 4px 8px;
                font-size: 12px;
                background-color: #e53e3e;
                color: white;
                border: none;
                border-radius: 3px;
                cursor: pointer;
            }
            #clearSessionBtn:hover {
                background-color: #c53030;
            }
        </style>
    </head>
    <body>
        <h1>聊天吧（简陋版）</h1>
        
        <div id="sessionInfo" class="session-info" style="display: none;">
            <span>当前会话ID: <span id="sessionIdDisplay" class="session-id"></span></span>
            <button id="clearSessionBtn" type="button">清除会话</button>
        </div>
        
        <div id="messages"></div>
        <form id="messageForm">
            <input type="text" id="messageInput" placeholder="输入消息..." required>
            <div class="controls">
                <div class="form-group">
                    <label for="modelSelected">模型:</label>
                    <select id="modelSelected">
                        <option value="qwen3-max">qwen3-max</option>
                        <option value="qwen-max">qwen-max</option>
                        <option value="qwen-plus">qwen-plus</option>
                        <option value="qwen-turbo">qwen-turbo</option>
                        <option value="qwen-flash">qwen-flash</option>
                        <option value="deepseek-v3.1">deepseek-v3.1</option>
                    </select>
                </div>
                <div class="form-group">
                    <label><input type="checkbox" id="streamChecked" checked>流式输出</label>
                </div>
                <button type="submit">发送</button>
                <button type="button" onclick="clearChat()">清空</button>
            </div>
        </form>
    
        <script>
            // 获取client_id
            let clientId = localStorage.getItem('clientId');
            if (!clientId) {
                clientId = 'client_' + Date.now();
                localStorage.setItem('clientId', clientId);
            }
            
            // 获取session_id
            let session_id = sessionStorage.getItem('session_id') || null;
            if (session_id) {
                document.getElementById('sessionIdDisplay').textContent = session_id;
                document.getElementById('sessionInfo').style.display = 'flex';
            }
            
            // 建立WebSocket连接
            const ws = new WebSocket(`ws://127.0.0.1:8011/ws/chat/${clientId}`);
            const messages = document.getElementById('messages');
            const messageInput = document.getElementById('messageInput');
            const streamChecked = document.getElementById('streamChecked');
            const modelSelected = document.getElementById('modelSelected');
            let currentStreamDiv = null;
    
            ws.onmessage = function(event) {
                const data = JSON.parse(event.data);
    
                if (data.status === 'running') {
                    if (streamChecked.checked === true) {
                        if (!currentStreamDiv) {
                            currentStreamDiv = addMessage('', 'assistant');
                        }
                        currentStreamDiv.textContent += data.message;
                    } else {
                        addMessage(data.message, 'assistant');
                    }
                } else if (data.status === 'completed' || data.status === 'cancelled') {
                    currentStreamDiv = null;
                } else if (data.status === 'error') {
                    addMessage('错误: ' + data.message, 'system');
                }
    
                messages.scrollTop = messages.scrollHeight;
                
                // 保存session_id到sessionStorage
                if (data.session_id) {
                    session_id = data.session_id;
                    sessionStorage.setItem('session_id', session_id);
                    
                    // 更新显示的session_id
                    document.getElementById('sessionIdDisplay').textContent = session_id;
                    document.getElementById('sessionInfo').style.display = 'flex';
                }
            };
    
            ws.onclose = function() {
                addMessage('连接已断开', 'system');
            };
    
            ws.onerror = function(error) {
                addMessage('连接错误: ' + error, 'system');
            };
    
            document.getElementById('messageForm').onsubmit = function(e) {
                e.preventDefault();
                const message = messageInput.value;
                if (message) {
                    addMessage(message, 'user');
                    
                    // 请求数据
                    const sendData = {
                        message: message,
                        stream: streamChecked.checked,
                        model: modelSelected.value
                    };
                    
                    // 增加session_id
                    if (session_id) {
                        sendData.session_id = session_id;
                    }
                    
                    ws.send(JSON.stringify(sendData));
                    messageInput.value = '';
                }
            };
    
            // 清除会话ID
            document.getElementById('clearSessionBtn').onclick = function() {
                session_id = null;
                sessionStorage.removeItem('session_id');
                document.getElementById('sessionInfo').style.display = 'none';
                clearChat()
            };
    
            function addMessage(message, type) {
                const div = document.createElement('div');
                div.className = 'message ' + type;
                div.textContent = message;
                messages.appendChild(div);
                return div;
            }
    
            function clearChat() {
                messages.innerHTML = '';
                currentStreamDiv = null;
            }
            
            // 显示当前session_id状态
            window.addEventListener('load', function() {
                if (session_id) {
                    addMessage(`已恢复会话: ${session_id}`, 'system');
                } else {
                    addMessage('开始新会话', 'system');
                }
            });
        </script>
    </body>
    </html>
    """)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=8011)

聊天吧: https://chat.wangzan.cc