Python基于FastAPI和WebSocket實現(xiàn)實時聊天應(yīng)用

更新時間：2025年04月30日 10:14:32 作者：老大白菜

這篇文章主要為大家詳細(xì)介紹了Python如何基于FastAPI和WebSocket實現(xiàn)實時聊天應(yīng)用,文中的示例代碼講解詳細(xì),感興趣的小伙伴可以跟隨小編一起學(xué)習(xí)一下

項目簡介

這是一個基于 FastAPI 和 WebSocket 實現(xiàn)的實時聊天應(yīng)用，支持一對一聊天、離線消息存儲等功能。

技術(shù)棧

后端：FastAPI (Python)

前端：HTML、JavaScript、CSS

通信：WebSocket

認(rèn)證：簡單的 token 認(rèn)證

項目結(jié)構(gòu)

├── main.py # 后端主程序
└── templates/ # 前端模板目錄
└── chat.html # 聊天頁面

詳細(xì)代碼實現(xiàn)

1. 后端實現(xiàn) (main.py)

from fastapi import FastAPI, WebSocket, Depends, WebSocketDisconnect
from typing import Dict
import json
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()
app.mount("/templates", StaticFiles(directory="templates"), name="templates")

class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.offline_messages: Dict[str, list] = {}  # 離線消息存儲

    async def connect(self, user: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[user] = websocket

    async def disconnect(self, user: str, websocket: WebSocket):
        if user in self.active_connections:
            del self.active_connections[user]

    async def send_personal_message(self, message: str, to_user: str):
        if to_user in self.active_connections:
            await self.active_connections[to_user].send_text(message)
            return {"success": True}
        else:
            # 存儲離線消息
            if to_user not in self.offline_messages:
                self.offline_messages[to_user] = []
            self.offline_messages[to_user].append(message)
            return {"success": False, "reason": "offline", "stored": True}

    async def get_offline_messages(self, user: str):
        if user in self.offline_messages:
            messages = self.offline_messages[user]
            del self.offline_messages[user]  # 取出后刪除
            return messages
        return []

# Token 認(rèn)證
async def get_cookie_or_token(token: str = None):
    if not token:
        from fastapi import HTTPException
        raise HTTPException(status_code=401, detail="未授權(quán)")
    return token

# 初始化連接管理器
ws_manager = ConnectionManager()

@app.websocket("/ws/{user}")
async def websocket_one_to_one(
        websocket: WebSocket,
        user: str,
        cookie_or_token: str = Depends(get_cookie_or_token)
):
    await ws_manager.connect(user, websocket)
    
    # 發(fā)送離線消息
    offline_msgs = await ws_manager.get_offline_messages(user)
    for msg in offline_msgs:
        await websocket.send_text(msg)
    
    try:
        while True:
            data = await websocket.receive_text()
            message_data = json.loads(data)
            
            # 防止自己給自己發(fā)消息
            if message_data["to_user"] == user:
                error_msg = {
                    "type": "error",
                    "content": "不能發(fā)送消息給自己"
                }
                await websocket.send_text(json.dumps(error_msg))
                continue
                
            response_message = {
                "from": user,
                "content": message_data["content"],
                "timestamp": message_data.get("timestamp"),
                "type": "message"
            }
            
            # 發(fā)送消息
            result = await ws_manager.send_personal_message(
                message=json.dumps(response_message),
                to_user=message_data["to_user"]
            )
            
            # 處理離線消息
            if not result["success"]:
                if result.get("stored"):
                    success_msg = {
                        "type": "info",
                        "content": f"消息已保存，將在用戶 {message_data['to_user']} 上線時送達(dá)"
                    }
                    await websocket.send_text(json.dumps(success_msg))
    except WebSocketDisconnect:
        await ws_manager.disconnect(user, websocket)
    except Exception as e:
        print(f"WebSocket 錯誤: {e}")
        await ws_manager.disconnect(user, websocket)

@app.get("/", response_class=FileResponse)
async def root():
    return "templates/chat.html"

2. 前端實現(xiàn) (templates/chat.html)

<!DOCTYPE html>
<html>
<head>
    <title>WebSocket Chat</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
        }
        #loginSection {
            text-align: center;
            margin-top: 100px;
        }
        #chatSection {
            display: none;
        }
        #messages {
            list-style-type: none;
            padding: 0;
            height: 400px;
            overflow-y: auto;
            border: 1px solid #ccc;
            margin-bottom: 20px;
            padding: 10px;
        }
        .message {
            margin: 10px 0;
            padding: 10px;
            border-radius: 5px;
            background-color: #f0f0f0;
        }
        .input-group {
            margin-bottom: 10px;
        }
        input[type="text"] {
            padding: 8px;
            margin-right: 10px;
            width: 200px;
        }
        button {
            padding: 8px 15px;
            background-color: #4CAF50;
            color: white;
            border: none;
            border-radius: 4px;
            cursor: pointer;
        }
        button:hover {
            background-color: #45a049;
        }
        #logoutBtn {
            background-color: #f44336;
        }
        #logoutBtn:hover {
            background-color: #da190b;
        }
        .error {
            color: red;
            margin: 10px 0;
        }
        .info {
            color: blue;
            margin: 10px 0;
        }
    </style>
</head>
<body>
    <!-- 登錄部分 -->
    <div id="loginSection">
        <h1>WebSocket 聊天</h1>
        <div class="input-group">
            <input type="text" id="loginUsername" placeholder="輸入用戶名" autocomplete="off"/>
            <button onclick="login()">登錄</button>
        </div>
        <div id="loginError" class="error"></div>
    </div>

    <!-- 聊天部分 -->
    <div id="chatSection">
        <h1>WebSocket 聊天</h1>
        <div id="currentUser"></div>
        <form action="" onsubmit="sendMessage(event)">
            <div class="input-group">
                <input type="text" id="messageText" placeholder="輸入消息" autocomplete="off"/>
                <input type="text" id="username" placeholder="接收者用戶名" autocomplete="off"/>
                <button type="submit">發(fā)送</button>
            </div>
        </form>
        <button id="logoutBtn" onclick="logout()">退出</button>
        <ul id='messages'></ul>
    </div>

    <script>
        let ws = null;

        function checkLogin() {
            const token = localStorage.getItem('token');
            if (token) {
                document.getElementById('loginSection').style.display = 'none';
                document.getElementById('chatSection').style.display = 'block';
                document.getElementById('currentUser').textContent = `當(dāng)前用戶: ${token}`;
                connectWebSocket();
            } else {
                document.getElementById('loginSection').style.display = 'block';
                document.getElementById('chatSection').style.display = 'none';
            }
        }

        function login() {
            const username = document.getElementById('loginUsername').value.trim();
            if (!username) {
                document.getElementById('loginError').textContent = '請輸入用戶名';
                return;
            }
            localStorage.setItem('token', username);
            checkLogin();
        }

        function logout() {
            if (ws && ws.readyState === WebSocket.OPEN) {
                ws.close();
            }
            localStorage.removeItem('token');
            document.getElementById('messages').innerHTML = '';
            document.getElementById('messageText').value = '';
            document.getElementById('username').value = '';
            document.getElementById('loginSection').style.display = 'block';
            document.getElementById('chatSection').style.display = 'none';
            document.getElementById('loginUsername').value = '';
            document.getElementById('currentUser').textContent = '';
        }

        function connectWebSocket() {
            const token = localStorage.getItem('token');
            ws = new WebSocket(`ws://localhost:8000/ws/${token}?token=${token}`);
            
            ws.onmessage = function(event) {
                const messages = document.getElementById('messages');
                const message = document.createElement('li');
                message.className = 'message';
                
                const data = JSON.parse(event.data);
                
                if (data.type === 'error') {
                    message.style.color = 'red';
                    message.textContent = data.content;
                } else if (data.type === 'info') {
                    message.style.color = 'blue';
                    message.textContent = data.content;
                } else {
                    message.textContent = `${data.from}: ${data.content}`;
                }
                
                messages.appendChild(message);
                messages.scrollTop = messages.scrollHeight;
            };

            ws.onerror = function(error) {
                console.error('WebSocket 錯誤:', error);
            };

            ws.onclose = function() {
                console.log('WebSocket 連接已關(guān)閉');
            };
        }

        function sendMessage(event) {
            event.preventDefault();
            
            const messageInput = document.getElementById("messageText");
            const usernameInput = document.getElementById("username");
            
            if (!messageInput.value.trim() || !usernameInput.value.trim()) {
                return;
            }
            
            const messageData = {
                content: messageInput.value,
                to_user: usernameInput.value,
                timestamp: new Date().toISOString()
            };

            ws.send(JSON.stringify(messageData));
            
            const messages = document.getElementById('messages');
            const message = document.createElement('li');
            message.className = 'message';
            message.style.backgroundColor = '#e3f2fd';
            message.textContent = `我: ${messageInput.value}`;
            messages.appendChild(message);
            messages.scrollTop = messages.scrollHeight;
            
            messageInput.value = '';
        }

        checkLogin();
    </script>
</body>
</html>