第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 -
在浏览器窗口输入: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)。你将学习如何将本地开发的应用部署到阿里云,包括容器化、数据库配置、自动化部署等内容。
希望这篇博客对你有所帮助!如果需要进一步扩展或调整内容,请随时告诉我。