使用Django-Channels實現websocket通訊+大模型對話

程序设计实验室發表於2024-08-14

前言

最近一直在做這個大模型專案,我選了 Django 作為框架(現在很多大模型應用都用的 FastAPI,不過我已經用習慣 Django 了)

之前使用 AspNetCore 作為後端的時候,我先後嘗試了 Blazor Server,WebAPI SSE(Server Sent Event)等方案來實現大模型對話,目前好像 SSE 是用得比較多的,ChatGPT 也是用的這個。

自從進入 3.x 時代,Django 開始支援非同步程式設計了,所以也能實現 SSE (不過我在測試中發現有點折騰),最終還是決定使用 WebSocket 來實現,Django 生態裡的這套東西就是 channels 了。(話說 FastAPI 好像可以直接支援 SSE 和 WebSocket )

先看效果

放一張聊天介面

聊天介面

我還做了個對話歷史頁面

對話歷史頁面

關於 django-channels

搬運一下官網的介紹,https://channels.readthedocs.io/en/latest/introduction.html

Channels wraps Django’s native asynchronous view support, allowing Django projects to handle not only HTTP, but protocols that require long-running connections too - WebSockets, MQTT, chatbots, amateur radio, and more.

OK,Channels 將 Django 從一個純同步的 HTTP 框架擴充套件到可以處理非同步協議如 WebSocket。原本 Django 是基於 WSGI 的,channels 使用基於 ASGI 的 daphne 伺服器,而且不止能使用 WebSocket ,還能支援 HTTP/2 等新的技術。

幾個相關的概念:

  • Channels: 持久的連線,如 WebSocket,可用於實時資料傳輸。
  • Consumers: 處理輸入事件的非同步功能,類似於 Django 的檢視,但專為非同步操作設計。
  • Routing: 類似於 Django 的 URL 路由系統,Channels 使用 routing 來決定如何分發給定的 WebSocket 連線或訊息到相應的 Consumer。

與傳統的 Django 請求處理相比,Django Channels 允許開發者使用非同步程式設計模式,這對於處理長時間執行的連線或需要大量併發連線的應用尤其有利。這種架構上的變化帶來了更高的效能和更好的使用者體驗,使 Django 能夠更好地適應現代網際網路應用的需求。

透過引入 Channels, Django 不再只是一個請求/響應式的 Web 框架,而是變成了一個真正意義上能夠處理多種網路協議和長時間連線的全功能框架。這使得 Django 開發者可以在不離開熟悉的環境的情況下,開發出更加豐富和動態的應用。

使用場景

先介紹下使用場景

這個 demo 專案的後端使用 StarAI 和 LangChain 呼叫 LLM 獲取回答,然後透過 WebSocket 與前端通訊,前端我選了 React + Tailwind

安裝

以 DjangoStarter 專案為例(使用 pdm 作為包管理器)

pdm add channels[daphne]

然後修改 src/config/settings/components/common.py

把 daphne 新增到註冊Apps裡,注意要放在最前面

# 應用定義
INSTALLED_APPS: Tuple[str, ...] = (
    'daphne',
)

之後使用 runserver 時,daphne 會代替 Django 內建的伺服器執行

接著,channels 還需要一個 channel layer,這可以讓多個消費者例項相互通訊以及與 Django 的其他部分通訊,這個 layer 可以選 Redis

pdm add channels_redis

配置

ASGI

OK,接下來得修改一下 src/config/asgi.py

import os

from channels.auth import AuthMiddlewareStack
from channels.routing import ProtocolTypeRouter, URLRouter
from channels.security.websocket import AllowedHostsOriginValidator
from django.core.asgi import get_asgi_application

os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'config.settings')

django_asgi_app = get_asgi_application()

from apps.chat.routing import websocket_urlpatterns

application = ProtocolTypeRouter({
    "http": django_asgi_app,
    # Just HTTP for now. (We can add other protocols later.)
    "websocket": AllowedHostsOriginValidator(
        AuthMiddlewareStack(URLRouter(websocket_urlpatterns))
    ),
})

除了官網文件說的配置之外,我這裡已經把聊天應用的路由加進來了

因為這個 demo 專案只有這一個 app 使用了 WebSocket ,所以這裡直接把 chat 裡的 routing 作為 root URLRouter

如果有多個 WebSocket app 可以按需修改

channel layer

修改 src/config/settings/components/channels.py 檔案

from config.settings.components.common import DOCKER

CHANNEL_LAYERS = {
    "default": {
        "BACKEND": "channels_redis.core.RedisChannelLayer",
        "CONFIG": {
            "hosts": [("redis" if DOCKER else "127.0.0.1", 6379)],
        },
    },
}

其他的配置就不贅述了,DjangoStarter 裡都配置好了

編寫後端程式碼

channels 的使用很簡單,正如前面的介紹所說,我們只需要完成 consumer 的邏輯程式碼,然後配置一下 routing 就行。

那麼開始吧

consumer

建立 src/apps/chat/consumers.py 檔案

身份認證、使用大模型生成回覆、聊天記錄等功能都在這了

先上程式碼,等下介紹

import asyncio
import json

from asgiref.sync import async_to_sync
from channels.db import database_sync_to_async
from channels.generic.websocket import WebsocketConsumer, AsyncWebsocketConsumer

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

from .models import Conversation, Message

class LlmConsumer(AsyncWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(args, kwargs)
        self.chat_id = None
        self.conversation = None

    @database_sync_to_async
    def get_conversation(self, pk, user):
        obj, _ = Conversation.objects.get_or_create(id=pk, user=user)
        return obj

    @database_sync_to_async
    def add_message(self, role: str, content: str):
        return Message.objects.create(
            conversation=self.conversation,
            role=role,
            content=content,
        )

    async def connect(self):
        self.chat_id = self.scope["url_route"]["kwargs"]["chat_id"]
        await self.accept()

        # 檢查使用者是否已登入
        user = self.scope["user"]
        if not user.is_authenticated:
            await self.close(code=4001, reason='Authentication required. Please log in again.')
            return
        else:
            self.conversation = await self.get_conversation(self.chat_id, user)

    async def disconnect(self, close_code):
        ...

    async def receive(self, text_data=None, bytes_data=None):
        text_data_json: dict = json.loads(text_data)
        history: list = text_data_json.get("history")

        user = self.scope["user"]
        if not user.is_authenticated:
            reason = 'Authentication required. Please log in again.'
            await self.send(text_data=json.dumps({"message": reason}))
            await asyncio.sleep(0.1)
            await self.close(code=4001, reason=reason)
            return

        await self.add_message(**history[-1])

        llm = ChatOpenAI(model="gpt4")
        resp = llm.stream(
            [
                *[
                    HumanMessage(content=e['content']) if e['role'] == 'user' else AIMessage(content=e['content'])
                    for e in history
                ],
            ]
        )

        message_chunks = []
        for chunk in resp:
            message_chunks.append(chunk.content)
            await self.send(text_data=json.dumps({"message": ''.join(message_chunks)}))
            await asyncio.sleep(0.1)

        await self.add_message('ai', ''.join(message_chunks))

工作流程

  • 前端發起ws連線之後,執行 connect 方法的程式碼,身份認證沒問題的話就呼叫 self.accept() 方法接受連線
  • 接收到前端發來的資訊,會自動執行 receive 方法,處理之後呼叫 self.send 可以傳送資訊

身份認證

雖然在 asgi.py 裡已經配置了 AuthMiddlewareStack

但還是需要在程式碼裡自行處理認證邏輯

這個中介軟體的作用是把 header 裡的 authorization 資訊變成 consumer 裡的 self.scope["user"]

要點

  • consumer 有同步版本和非同步版本,這裡我使用了非同步版本 AsyncWebsocketConsumer
  • async consumer 裡訪問 Django ORM 需要使用 database_sync_to_async 裝飾器
  • connect 裡面,身份驗證失敗後呼叫 self.close 關閉連線,裡面的引數 code 不能和 HTTP Status Code 衝突,我一開始用的 401 ,結果前端接收時變成其他 code ,換成自定義的 4001 才可以接收到。但 reason 引數是一直接收不到的,不知道為啥,可能跟瀏覽器的 WebSocket 實現有關?
  • receive 方法裡,將大模型生成的內容使用流式輸出傳送給客戶端時,一定要在 for 迴圈里加上 await asyncio.sleep 等待一段時間,這是為了留出時間讓 WebSocket 傳送訊息給客戶端,不然會變成全部生成完再一次性發給客戶端,沒有流式輸出的效果。
  • receive 方法裡,接收到訊息後先判斷當前是否有登入(或者登入是否過期,表現和未登入一致),如果未登入則傳送資訊 "Authentication required. Please log in again." 給客戶端,然後再關閉連線。

routing

寫完了 consumer ,配置一下路由

編輯 src/apps/chat/routing.py 檔案

from django.urls import re_path

from . import consumers

websocket_urlpatterns = [
    re_path(r"ws/chat/demo/(?P<room_name>\w+)/$", consumers.ChatConsumer.as_asgi()),
    re_path(r"ws/chat/llm/(?P<chat_id>\w+)/$", consumers.LlmConsumer.as_asgi()),
]

這裡注意只能使用 re_path 不能 path (官方文件說的)

客戶端開發

OK,後端部分到這就搞定了,接下來寫一下客戶端的程式碼

我選擇了 React 來實現客戶端

無關程式碼就不放了,直接把關鍵的程式碼貼上來

function App() {
  const [prompt, setPrompt] = useState('')
  const [messages, setMessages] = useState([])
  const [connectState, setConnectState] = useState(3)
  const [conversation, setConversation] = useState()

  const chatSocket = useRef(null)
  const messagesRef = useRef(null)
  const reLoginModalRef = useRef(null)

  React.useEffect(() => {
    openWebSocket()
    getConversation()

    return () => {
      chatSocket.current.close();
    }
  }, [])

  React.useEffect(() => {
    // 自動滾動到訊息容器底部
    messagesRef.current.scrollIntoView({behavior: 'smooth'})
  }, [messages]);

  const getConversation = () => {
    // ... 省略獲取聊天記錄程式碼
  }

  const openWebSocket = () => {
    setConnectState(0)
    chatSocket.current = new WebSocket(`ws://${window.location.host}/ws/chat/llm/${ConversationId}/`)
    chatSocket.current.onopen = function (e) {
      console.log('WebSocket連線建立', e)
      setConnectState(chatSocket.current.readyState)
    };

    chatSocket.current.onmessage = function (e) {
      const data = JSON.parse(e.data)
      setMessages(prevMessages => {
        if (prevMessages[prevMessages.length - 1].role === 'ai')
          return [
            ...prevMessages.slice(0, -1), {role: 'ai', content: data.message}
          ]
        else
          return [
            ...prevMessages, {role: 'ai', content: data.message}
          ]
      })
    };

    chatSocket.current.onclose = function (e) {
      console.error('WebSocket 連結斷開。Chat socket closed unexpectedly.', e)
      setConnectState(chatSocket.current.readyState)
      if (e.code === 4001) {
        // 顯示重新登入對話方塊
        new Modal(reLoginModalRef.current, options).show()
      }
    };
  }

  const sendMessage = () => {
    if (prompt.length === 0) return

    const newMessages = [...messages, {
      role: 'user', content: prompt
    }]
    setMessages(newMessages)
    setPrompt('')
    chatSocket.current.send(JSON.stringify({
      'history': newMessages
    }))
  }
}

前端程式碼沒啥好說的,簡簡單單

我用了 messagesRef.current.scrollIntoView({behavior: 'smooth'}) 來實現訊息自動滑動到底部,之前用 Blazor 開發 AIHub 的時候好像是用了其他實現,我記得不是這個

不過這個方法挺絲滑的

WebSocket 地址直接寫在這裡面感覺沒那麼優雅,而且後面部署後得改成 wss:// 也麻煩,得研究一下有沒有更優雅的實現。

小結

這個只是簡單的demo,實際上生產還得考慮很多問題,本文就是為 channels 的應用開了個頭,後續有新的研究成果會持續更新部落格~

相關文章