FastAPI: websocket的用法及举例

1. Websocket

1.1 Websocket介绍

WebSocket 是一种在单个TCP连接上进行全双工通信的协议，允许客户端和服务器之间相互发送数据，而不需要像传统的HTTP请求-响应模型那样频繁建立和断开连接。
全双工通信(Full-Duplex Communication)是一种通信模式，允许通信双方同时发送和接收数据。换句话说，数据可以同时从两端双向传输，而不会相互阻塞或干扰。

1.2 FastAPI中的Websocket

FastAPI提供了对WebSocket的原生支持，允许你轻松构建高效的实时应用，如聊天室、实时数据更新等。

1.2.1 装饰器

FastAPI中与WebSocket相关的主要装饰器为 @app.websocket。该装饰器的作用和参数如下：

作用：将一个路径(如/ws)与一个处理WebSocket请求的函数关联。当客户端通过WebSocket连接该路径时，FastAPI会调用该函数处理连接和通信。
参数：它接受的参数与其他路由装饰器相同，主要是路径(URL)，可选地也能设置依赖项、权限等。

代码举例如下(当客户端通过WebSocket连接/ws路径时，FastAPI将执行下面的websocket_endpoint函数)：

from fastapi import FastAPI, WebSocket
app = FastAPI()
# 定义一个 WebSocket 路由
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()  # 接受 WebSocket 连接
    while True:
        data = await websocket.receive_text()  # 接收来自客户端的消息
        await websocket.send_text(f"Message text was: {data}")  # 回复消息给客户端

1.2.2 websocket相关方法

FastAPI提供了处理WebSocket各种事件的方法，包括接受消息、发送消息、关闭连接等。具体如下：

websocket.accept：接受WebSocket连接请求。
websocket.receive_text：接收客户端发来的文本消息。
websocket.send_text：向客户端发送文本消息。
websocket.close：关闭WebSocket连接。

2. 构建对话机器人

这里我们用FastAPI和React构建一个聊天机器人的聊天界面。这里关于机器人的后端处理逻辑这里不做详细介绍，而前端部分主要介绍App.tsx和ChatPag.tsx内容，不介绍CSS部分。具体代码如下：
React中App.tsx代码如下：

import './App.css';
import ChatPage from './components/ChatPage';

function App() {
  return (
    <div className="App">
        <div className="header">
          <div className="header-logo">
            <img src="https://cdn.builder.io/api/v1/image/assets/TEMP/b0db057162d379f22892cd5ae4d13c509717e0a81da39be3f65cb94e15556ed7?apiKey=0682bce60b3549f085131079f1bf89f0&&apiKey=0682bce60b3549f085131079f1bf89f0" alt="Chainlit" /> 
            <div className="header-title">SmartRecommend服务推荐助手</div>
          </div>
        </div>
        <div className='body-container'>
          <div className="main">
            <div className='chatpage'>
              <ChatPage />
            </div>
          </div>
        </div>
      </div>
  );
}
export default App;

React中ChatPage.tsx代码如下:

import "./ChatPage.css";
import { useEffect, useState} from "react";
import { nanoid } from 'nanoid';

interface Message{
    id:string,
    name:string,
    type:string,
    output:string,
    createdAt:number|string,
}
function ChatPage() {
    const [inputValue, setInputValue] = useState("");
    const [messages,setMessages] = useState<Message[]>([]);
    const [socket,setSocket] = useState<WebSocket|null>(null);
    
    useEffect(() => {
        const ws = new WebSocket("ws://localhost:8000/ws/chat");
        ws.onopen = () => {
          console.log("websocket链接已建立!");
        };
        ws.onmessage = (event) => {
          const message = JSON.parse(event.data);
          setMessages((prevMessages) => [...prevMessages, message]);
        };
        ws.onerror = (error) => {
            console.log('WebSocket错误:', error);
        };
        ws.onclose=()=>{
            console.log("websocket链接已关闭!");
        }
        setSocket(ws);
        return () => {
            ws.close();
        }
    },[]);

    const handleSendMessage = () => {
        const content = inputValue.trim();
        if (content) {
          const message: Message={
            id: nanoid(),
            name: "User",
            type: "user_message",
            output: content,
            createdAt: Date.now(),
          };
        setMessages((prevMessages) => [...prevMessages, message]);
        socket?.send(JSON.stringify(message));
        }
        setInputValue("");
      };

    const renderMessage = (message:Message,index:number) => {
        const dateOptions: Intl.DateTimeFormatOptions = {
            hour: "2-digit",
            minute: "2-digit",
          };
          const date = new Date(message.createdAt).toLocaleTimeString(
            undefined,
            dateOptions
          );
          if(message.type === "user_message") {
            return (
                <div key={message.id} className="chat-box-user">
                  <div className="user-avatar">U</div>
                  <div className="bot-user-content">
                    <div className="user-icon">
                      <div className="bot-user-name">{message.name}</div>
                      <div className="bot-user-time">{date}</div>
                    </div>
                    <div className="user-chat-message">{message.output}</div>
                  </div>
                </div>
            );
          } else {
              return (
                <div key={message.id} className="chat-box-bot">
                  <div className="bot-avatar">B</div>
                  <div className="bot-user-content">
                    <div className="bot-icon">
                        <div className="bot-user-name">{message.name}</div>
                        <div className="bot-user-time">{date}</div>
                    </div>
                    <div className="bot-chat-message">{message.output}</div>
                  </div>
                </div>
              );
          };
        };
    return (
        <div className="chat-container">
        <div className="chat-box">
            {messages.map(renderMessage)}
        </div>
        <div className="fixed-bottom">
          <input className="fixed-bottom-input" 
                type="text"
                value={inputValue}
                placeholder="你可以输入“你好”唤醒服务"
                onChange={(e) => setInputValue(e.target.value)}
                onKeyUp={(e) => {
                  if (e.key === "Enter") {
                    handleSendMessage();
                  }
                }}>
          </input>
          <button onClick={handleSendMessage} className="button" type="submit">Send</button>
        </div>
      </div>
    ); 
}
export default ChatPage;

后端FastAPI代码：

from fastapi import FastAPI, WebSocket,HTTPException
import uvicorn
from fastapi.middleware.cors import CORSMiddleware
from typing import List
import json
import datetime
from nanoid import generate
import httpx

app=FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
clients: List[WebSocket] = []

RASA_API_URL="http://localhost:5005/webhooks/rest/webhook"

@app.websocket("/ws/chat")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    clients.append(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            for client in clients:
                text={"id": generate(),
                      "name":"Bot","type":"bot_message",
                      "output":json.loads(data)["output"],
                      "createdAt":int(datetime.datetime.now().timestamp()*1000)}
                text=json.dumps(text)
                await client.send_text(text)
    except Exception as e:
        print(e)
        clients.remove(websocket)
        
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)