Потоковая передача WebSocket Не работает в Django + Next.js — Получаем только первое и последнее сообщение

Я создаю приложение для чата с искусственным интеллектом, используя Django (каналы) для серверной части и Next.js для интерфейса. Цель состоит в том, чтобы передавать сгенерированные ИИ ответы по частям через соединение WebSocket - аналогично тому, как работает REST API с ReadableStream. Однако на интерфейсе отображается только первый фрагмент, а затем окончательное завершенное сообщение, а не промежуточные потоковые фрагменты.

Вот упрощенная версия моей настройки:

Серверная часть (Django WebSocket Consumer): Я использую асинхронный генератор (handle_chat_request) для создания фрагментов сообщений.

async def receive(self, text_data):
    logging.info(f"WebSocket message received: {text_data}")

    try:
        data = json.loads(text_data)
        message_type = data.get("type")

        if message_type == "message":
            jwt_token = data.get("token")
            if not jwt_token:
                await self.send(json.dumps({"type": "error", "message": "Missing Authorization token."}))
                await self.close(code=4001)
                return

            user_message = data.get("content")
            user_id = data.get("user_id")

            combined_response = ""

            # Stream response chunks
            async for chunk in handle_chat_request(user_message, user_id, jwt_token):
                combined_response += chunk
                await self.send(json.dumps({"type": "ai_response_chunk", "content": chunk}))
                await asyncio.sleep(0)  # Yield control to allow chunk sending

            # Send final complete message
            await self.send(json.dumps({"type": "ai_response_complete", "content": combined_response}))
        
    except Exception as e:
        logging.error(f"WebSocket error: {e}")
        await self.send(json.dumps({"type": "error", "message": "An error occurred."}))

Интерфейс (Next.js Клиент WebSocket): Я прослушиваю входящие фрагменты и пытаюсь добавить их к последнему сообщению бота:

const handleWebSocketMessage = (event) => {
  try {
    const data = JSON.parse(event.data);

    switch (data.type) {
      case "ai_response_chunk":
        setMessages((messages) => {
          const updatedMessages = [...messages];
          const lastMessageIndex = updatedMessages.findLastIndex(
            (m) => m.role === "bot"
          );
          
          if (lastMessageIndex !== -1) {
            updatedMessages[lastMessageIndex] = {
              ...updatedMessages[lastMessageIndex],
              content: updatedMessages[lastMessageIndex].content + data.content,
            };
          } else {
            updatedMessages.push({ role: "bot", content: data.content });
          }

          return updatedMessages;
        });
        break;

      case "ai_response_complete":
        //some other functionality
        break;

      case "error":
        setError(data.message);
        setLoading(false);
        break;
    }
  } catch (error) {
    console.error("Error parsing WebSocket message:", error);
    setLoading(false);
  }
};

Проблема: "ai_response_chunk" работает так, как и ожидалось от пользовательского интерфейса, отображает только первый фрагмент, а затем конечное содержимое сообщения, когда оно будет завершено, но я не вижу дополнительных обновлений, как в случае с реализацией REST API.

Реализация Rest:

  fetch(`/api/chat`, {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
          },
          body: JSON.stringify({ message: userMessage }),
        }).then(async (response) => {
          if (!response.ok) throw new Error("Chat API failed");
          if (!response.body) throw new Error("ReadableStream not supported");

          const reader = response.body.getReader();
          const decoder = new TextDecoder();
          let done = false;
          let botMessage = "";

          // Add bot message placeholder
          setMessages((messages) => [
            ...messages,
            { content: "", role: "bot" },
          ]);

          while (!done) {
            const { value, done: doneReading } = await reader.read();
            done = doneReading;
            const chunk = decoder.decode(value, { stream: true });
            botMessage += chunk;

            setMessages((messages) => {
              const updatedMessages = [...messages];
              const lastMessageIndex = updatedMessages.findLastIndex(
                (m) => m.role === "bot"
              );
              if (lastMessageIndex !== -1) {
                updatedMessages[lastMessageIndex] = {
                  ...updatedMessages[lastMessageIndex],
                  content: updatedMessages[lastMessageIndex].content + chunk,
                };
              }
              return updatedMessages;
            });
          }
    ```

for chunk in self.sync_stream_chat_response(user_message, user_id, jwt_token):
    chunk_count += 1
    print(f"Sending chunk #{chunk_count} to WebSocket")
    await self.send(json.dumps({"type": "ai_response_chunk", "content": chunk}))
    
    # Add a small delay (e.g., 50ms)
    await asyncio.sleep(0.05)

await self.send(json.dumps({"type": "ai_response_complete"}))