《Python实战进阶》No 7: 一个AI大模型聊天室的构建-基于WebSocket 实时通信开发实战

第7集： 一个AI大模型聊天室的构建-基于WebSocket 实时通信开发实战

在现代 Web 开发中，实时通信已经成为许多应用的核心需求。无论是聊天应用、股票行情推送，还是多人协作工具，WebSocket 都是实现高效实时通信的最佳选择之一。本集将一个AI大模型聊天室的实战案例，带你深入了解 WebSocket 的基础知识、Python 实现方式以及实际应用场景。

以下截图为 AI助手多人聊天室 页面，完整代码和搭建说明在第4节中提供：

---

1. WebSocket 基础知识

1.1 WebSocket 协议与 HTTP 的区别

WebSocket 是一种全双工通信协议，允许客户端和服务器之间建立持久连接，从而实现低延迟的数据传输。与传统的 HTTP 请求-响应模式相比，WebSocket 有以下特点：

特性	HTTP	WebSocket
连接类型	短连接（每次请求后断开）	长连接（持久化）
数据传输方向	单向（客户端 -> 服务器）	双向（客户端 <-> 服务器）
性能	每次请求需要重新建立连接	一次握手后持续通信
应用场景	页面加载、表单提交等	实时聊天、股票行情推送等

1.2 WebSocket 的应用场景

WebSocket 的实时性和高效性使其适用于以下场景：

• 实时聊天：用户之间的即时消息传递。
• 股票行情推送：实时更新股票价格。
• 在线游戏：玩家之间的实时互动。
• 通知系统：订单状态更新、系统告警等。

1.3 WebSocket 的握手过程与数据帧格式

WebSocket 的通信从 HTTP 握手开始，握手完成后切换到 WebSocket 协议。以下是握手过程的简化流程：

1. 客户端发送一个带有 Upgrade: websocket 的 HTTP 请求。
2. 服务器响应 101 Switching Protocols，表示协议切换成功。
3. 双方通过 WebSocket 协议进行数据传输。

WebSocket 数据帧格式如下：

• FIN：标识是否为最后一帧。
• Opcode：操作码（如文本帧、二进制帧）。
• Payload：实际传输的数据。

---

2. Python 中实现 WebSocket

2.1 使用 websockets 库构建 WebSocket 服务器和客户端

websockets 是一个简单易用的 Python 库，用于实现 WebSocket 通信。以下是一个简单的示例：

Python版本：3.11.5, 建议新建一个虚拟环境目录 websocket **
** 安装依赖:

python -m venv websocket
pip install asyncio websockets

服务器端代码 server.py

import asyncio
import websockets

async def echo(websocket):
    # 注意这里移除了 path 参数
    async for message in websocket:
        print(f"Received: {message}")
        await websocket.send(f"Echo: {message}")

async def main():
    # 在异步函数内创建并启动服务器
    server = await websockets.serve(echo, "localhost", 8765)
    print("WebSocket server started at ws://localhost:8765")
    # 保持服务器运行
    await asyncio.Future()  # 这会一直运行，直到被取消

# 使用新的异步运行方式
if __name__ == "__main__":
    asyncio.run(main())

客户端代码 client.py

import asyncio
import websockets

async def hello():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        await websocket.send("Hello, WebSocket!")
        response = await websocket.recv()
        print(f"Received: {response}")

# 使用新的异步运行方式
if __name__ == "__main__":
    asyncio.run(hello())

分别在两个cmd启动同一个虚拟环境 scripts/activate ，然后先运行 python server.py
在另一个cmd运行 python client.py
运行上述代码后，客户端会向服务器发送消息，服务器会返回一个回显消息。
注意：要开两个cmd或shell窗口，一个运行服务，一个运行客户端

客户端

(websocket) D:\python_projects\websocket>python client.py
Received: Echo: Hello, WebSocket!

服务端

WebSocket server started at ws://localhost:8765
Received: Hello, WebSocket!

2.2 使用 Flask 或 Django 集成 WebSocket

对于 Flask 和 Django，可以分别使用 Flask-SocketIO 和 Django Channels 来集成 WebSocket。
以下是简单示例，不用上手代码，在第4节有完整的基于socket的AI聊天室代码

Flask-SocketIO 示例

from flask import Flask, render_template
from flask_socketio import SocketIO

app = Flask(__name__)
socketio = SocketIO(app)

@app.route("/")
def index():
    #return render_template("index.html")
    return "WebSocket Server is running!"

@socketio.on("message")
def handle_message(message):
    print(f"Received: {message}")
    socketio.emit("response", f"Echo: {message}")

if __name__ == "__main__":
    socketio.run(app, host="0.0.0.0", port=5000)

Django Channels 示例

1. 安装 channels 并配置 settings.py。
2. 编写消费者逻辑：

from channels.generic.websocket import AsyncWebsocketConsumer
import json

class ChatConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        await self.accept()

    async def disconnect(self, close_code):
        pass

    async def receive(self, text_data):
        text_data_json = json.loads(text_data)
        message = text_data_json["message"]
        await self.send(text_data=json.dumps({"message": f"Echo: {message}"}))

---

3. 性能优化与安全性

3.1 如何处理高并发连接

• 使用异步框架（如 asyncio 或 uvloop）提高性能。
• 部署负载均衡器（如 Nginx）分担流量压力。

3.2 WebSocket 的安全问题

• 防止 CSRF 攻击：验证 WebSocket 连接来源。
• 数据加密：使用 SSL/TLS 加密 WebSocket 通信。

3.3 使用 SSL/TLS 加密 WebSocket 通信

在生产环境中，建议使用 HTTPS 和 WSS（WebSocket Secure）。以下是一个启用 SSL 的示例：

ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ssl_context.load_cert_chain(certfile="cert.pem", keyfile="key.pem")

start_server = websockets.serve(echo, "localhost", 8765, ssl=ssl_context)

---

4. 实际案例

构建一个基于AI的Websocket聊天室

使用flask框架，基于本地部署的ollama提供大模型接口，构建一个用户与人工智能助手的聊天室。

主程序代码 websocket_demo.py :

from flask import Flask, render_template
from flask_socketio import SocketIO, emit
import requests
import json

app = Flask(__name__)
socketio = SocketIO(app)

# Ollama API 配置
OLLAMA_API_URL = "http://localhost:11434/api/generate"
MODEL_NAME = "glm4:latest"  # 可以根据你本地安装的模型修改

# 使用 Ollama 生成回复
def get_ollama_response(prompt):
    try:
        payload = {
            "model": MODEL_NAME,
            "prompt": prompt,
            "stream": False
        }
        
        response = requests.post(OLLAMA_API_URL, json=payload)
        
        if response.status_code == 200:
            result = response.json()
            return result.get("response", "抱歉，我无法生成回复。")
        else:
            print(f"Ollama API 错误: {response.status_code}")
            return "抱歉，连接大模型时出现错误。"
    except Exception as e:
        print(f"调用 Ollama 时出错: {str(e)}")
        return "抱歉，连接大模型时出现错误。"

# 路由：返回 HTML 页面
@app.route("/")
def index():
    return render_template("index.html")

# WebSocket 事件处理
@socketio.on("message")
def handle_message(data):
    print(f"Received message: {data}")
    
    user_message = data["message"]
    username = data["username"]
    
    # 发送用户原始消息
    emit("response", {"message": user_message, "username": username}, broadcast=True)
    
    # 获取 Ollama 的回复
    prompt = f"用户 {username} 说: {user_message}\n请简短回复:"
    ai_response = get_ollama_response(prompt)
    
    # 发送 AI 回复
    emit("response", {"message": ai_response, "username": "AI助手"}, broadcast=True)

if __name__ == "__main__":
    print("Starting Flask-SocketIO WebSocket server with Ollama integration...")
    socketio.run(app, host="0.0.0.0", port=5000)

在虚拟环境主目录下新建一个templates文件夹，放置以下index.html页面文件:

<!DOCTYPE html>
<html lang="zh">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>基于AI的 WebSocket 聊天室</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
        }
        #chat-box {
            height: 400px;
            border: 1px solid #ccc;
            padding: 10px;
            margin-bottom: 10px;
            overflow-y: auto;
        }
        .message {
            margin-bottom: 10px;
            padding: 8px;
            border-radius: 5px;
        }
        .user-message {
            background-color: #e6f7ff;
        }
        .system-message {
            background-color: #f0f0f0;
            font-style: italic;
        }
        .username {
            font-weight: bold;
            margin-right: 10px;
        }
        #message-form {
            display: flex;
        }
        #message-input {
            flex-grow: 1;
            padding: 8px;
            margin-right: 10px;
        }
    </style>
</head>
<body>
    <h1>WebSocket 聊天室</h1>
    
    <div id="username-container">
        <label for="username">请输入用户名：</label>
        <input type="text" id="username" value="用户">
        <button id="join-btn">加入聊天</button>
    </div>
    
    <div id="chat-container" style="display: none;">
        <div id="chat-box"></div>
        
        <form id="message-form">
            <input type="text" id="message-input" placeholder="输入消息..." autocomplete="off">
            <button type="submit">发送</button>
        </form>
    </div>

    <script src="https://cdnjs.cloudflare.com/ajax/libs/socket.io/4.0.1/socket.io.js"></script>
    <script>
        document.addEventListener('DOMContentLoaded', () => {
            const socket = io();
            const chatBox = document.getElementById('chat-box');
            const messageForm = document.getElementById('message-form');
            const messageInput = document.getElementById('message-input');
            const usernameInput = document.getElementById('username');
            const joinBtn = document.getElementById('join-btn');
            const usernameContainer = document.getElementById('username-container');
            const chatContainer = document.getElementById('chat-container');
            
            let username = '';
            
            // 加入聊天
            joinBtn.addEventListener('click', () => {
                username = usernameInput.value.trim() || '用户';
                usernameContainer.style.display = 'none';
                chatContainer.style.display = 'block';
                
                // 发送加入消息
                socket.emit('message', {
                    message: '加入了聊天室',
                    username: username
                });
            });
            
            // 发送消息
            messageForm.addEventListener('submit', (e) => {
                e.preventDefault();
                const message = messageInput.value.trim();
                
                if (message) {
                    socket.emit('message', {
                        message: message,
                        username: username
                    });
                    messageInput.value = '';
                }
            });
            
            // 接收消息
            socket.on('response', (data) => {
                const messageDiv = document.createElement('div');
                messageDiv.className = data.username === '系统' ? 'message system-message' : 'message user-message';
                
                const usernameSpan = document.createElement('span');
                usernameSpan.className = 'username';
                usernameSpan.textContent = data.username + ':';
                
                const messageContent = document.createElement('span');
                messageContent.textContent = data.message;
                
                messageDiv.appendChild(usernameSpan);
                messageDiv.appendChild(messageContent);
                chatBox.appendChild(messageDiv);
                
                // 自动滚动到底部
                chatBox.scrollTop = chatBox.scrollHeight;
            });
        });
    </script>
</body>
</html> 

运行主程序：

python websocket_demo.py

cmd窗口输出如下:

Starting Flask-SocketIO WebSocket server with Ollama integration...
 * Serving Flask app 'websocket_demo'
 * Debug mode: off
WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.   
 * Running on all addresses (0.0.0.0)
 * Running on http://127.0.0.1:5000
 * Running on http://192.168.31.52:5000
Press CTRL+C to quit
127.0.0.1 - - [27/Feb/2025 21:29:42] "GET /socket.io/?EIO=4&transport=polling&t=PL7cKLt HTTP/1.1" 200 -
127.0.0.1 - - [27/Feb/2025 21:29:42] "GET / HTTP/1.1" 200 -
127.0.0.1 - - [27/Feb/2025 21:29:42] "GET /socket.io/?EIO=4&transport=polling&t=PL7cKV2 HTTP/1.1" 200 -
127.0.0.1 - - [27/Feb/2025 21:29:42] "POST /socket.io/?EIO=4&transport=polling&t=PL7cKVc&sid=IYT6sUZFDo9kVtxHAAAB HTTP/1.1" 200 -
127.0.0.1 - - [27/Feb/2025 21:29:42] "GET /socket.io/?EIO=4&transport=polling&t=PL7cKVd&sid=IYT6sUZFDo9kVtxHAAAB HTTP/1.1" 200 -
127.0.0.1 - - [27/Feb/2025 21:29:42] "GET /socket.io/?EIO=4&transport=polling&t=PL7cKZE&sid=IYT6sUZFDo9kVtxHAAAB HTTP/1.1" 200 -

在浏览器窗口输入：localhost:5000

---

输入用户名 哪吒

随便聊一句：

聊天室支持多人聊天，新开一个浏览器窗口输入：localhost:5000
输入用户名：孙悟空，在所有用户窗口会同步更新所有聊天消息。

一个实际可用的多人AI聊天室就构建完成了，

用户可以根据需要修改大模型资源接口和prompt以达到定制效果。

5. 扩展方向

5.1 对比 WebSocket 与其他实时通信技术

• Server-Sent Events (SSE)：仅支持服务器到客户端的单向通信。
• 长轮询：模拟实时通信，但效率较低。

5.2 在微服务架构中使用 WebSocket

• 将 WebSocket 服务独立部署为一个微服务。
• 使用消息队列（如 RabbitMQ）解耦 WebSocket 服务与其他服务。

5.3 基于DEMO进行开发

• 增加多方通讯的管理，利用大模型进行聊天逻辑和内容管控。
• 完善聊天室的功能，比如表情包、图片和语音文件的上传。

---

下集预告

下一集我们将聚焦于 No8: 部署 Flask/Django 应用到云平台（Aliyun）。你将学习如何将本地开发的应用部署到阿里云，包括容器化、数据库配置、自动化部署等内容。

---

希望这篇博客对你有所帮助！如果需要进一步扩展或调整内容，请随时告诉我。