《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 -