FastAPI WebSocket 聊天室：用戶加入廣播 + 私聊實現

作者：Ss肥魚 2025-05-09 08:35:00

在實際聊天室中，我們通常希望有廣播系統消息和私聊功能，而FastAPI + WebSocket 支持這些非常輕松。

在實際聊天室中，我們通常希望有以下功能：

廣播系統消息：某個用戶加入/離開聊天室，通知所有人。
私聊功能：用戶可以給指定用戶發送私密消息。

FastAPI + WebSocket 支持這些非常輕松。

目錄結構

fastapi-chat/
├── main.py
└── chat_manager.py

依賴安裝

pip install fastapi uvicorn

(1) chat_manager.py：升級連接管理器，支持用戶身份

# chat_manager.py
from fastapi import WebSocket
from typing import Dict


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, username: str, websocket: WebSocket):
        await websocket.accept()
        self.active_connections[username] = websocket
        await self.broadcast(f"?? 用戶【{username}】加入了聊天室")

    def disconnect(self, username: str):
        if username in self.active_connections:
            del self.active_connections[username]

    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)

    async def broadcast(self, message: str):
        for ws in self.active_connections.values():
            await ws.send_text(message)

(2) main.py：WebSocket 支持用戶名、私聊

# main.py
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
from chat_manager import ConnectionManager

app = FastAPI()
manager = ConnectionManager()

@app.websocket("/ws/chat/")
async def websocket_endpoint(websocket: WebSocket, username: str = Query(...)):
    await manager.connect(username, websocket)
    try:
        while True:
            data = await websocket.receive_json()
            msg_type = data.get("type", "broadcast")
            content = data.get("message")
            to_user = data.get("to")

            if msg_type == "private" and to_user:
                await manager.send_personal_message(f"?? 私信【{username}】→【{to_user}】：{content}", to_user)
            else:
                await manager.broadcast(f"?? 【{username}】：{content}")
    except WebSocketDisconnect:
        manager.disconnect(username)
        await manager.broadcast(f"? 用戶【{username}】離開了聊天室")

HTML 前端示例（支持私聊發送）

<!-- index.html -->
<!DOCTYPE html>
<html>
<head><title>聊天</title></head>
<body>
  <h2>FastAPI 聊天室</h2>
  <input id="username" placeholder="用戶名" />
  <button onclick="connect()">連接</button>

  <div>
    <input id="to" placeholder="發給誰（留空為群發）" />
    <input id="message" placeholder="消息內容" />
    <button onclick="send()">發送</button>
  </div>

  <ul id="messages"></ul>

  <script>
    let ws = null;
    function connect() {
      const username = document.getElementById("username").value;
      ws = new WebSocket("ws://localhost:8000/ws/chat/?username=" + username);
      ws.onmessage = event => {
        const li = document.createElement("li");
        li.innerText = event.data;
        document.getElementById("messages").appendChild(li);
      };
    }

    function send() {
      const msg = document.getElementById("message").value;
      const to = document.getElementById("to").value;
      const type = to ? "private" : "broadcast";
      ws.send(JSON.stringify({ type, message: msg, to }));
    }
  </script>
</body>
</html>

功能回顧

新用戶連接 → 自動廣播 ?? 用戶加入
用戶斷開連接 → 自動廣播 ? 用戶離開
留空接收人 → 群發廣播消息
指定 to 字段 → 發送私聊

責任編輯：趙寧寧來源： Ssoul肥魚