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构建一个聊天机器人。这里关于机器人的后端处理逻辑这里不做详细介绍(这里不介绍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 httpxapp=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)