《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 协议。以下是握手过程的简化流程：

客户端发送一个带有 Upgrade: websocket 的 HTTP 请求。
服务器响应 101 Switching Protocols，表示协议切换成功。
双方通过 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 websocketsasync 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 websocketsasync 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 SocketIOapp = 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 示例

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

from channels.generic.websocket import AsyncWebsocketConsumer
import jsonclass ChatConsumer(AsyncWebsocketConsumer):async def connect(self):await self.accept()async def disconnect(self, close_code):passasync 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 jsonapp = 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 -