Real-time AI Features: WebSocket + LLM Streaming Trong React

HTTP request-response pattern không đủ cho real-time AI experiences. Khi bạn cần multiple users thấy AI responses cùng lúc, hoặc AI cần push updates liên tục, WebSocket + LLM streaming là answer. Bài này hướng dẫn build real-time AI features production-ready.

Architecture Overview

Stack: React frontend → WebSocket server (Node.js) → LLM API (streaming). WebSocket server acts as proxy — nhận messages từ clients, forward đến LLM, stream responses back cho all connected clients.

WebSocket Server với LLM Streaming

// server/ws-ai.ts
import { WebSocketServer, WebSocket } from "ws";
import { streamText } from "ai";
import { openai } from "@ai-sdk/openai";

const wss = new WebSocketServer({ port: 8080 });
const rooms = new Map<string, Set<WebSocket>>();

wss.on("connection", (ws) => {
  let currentRoom: string | null = null;

  ws.on("message", async (data) => {
    const msg = JSON.parse(data.toString());

    if (msg.type === "join") {
      currentRoom = msg.roomId;
      if (!rooms.has(msg.roomId)) rooms.set(msg.roomId, new Set());
      rooms.get(msg.roomId)!.add(ws);
      return;
    }

    if (msg.type === "ai-query" && currentRoom) {
      const room = rooms.get(currentRoom);
      if (!room) return;

      // Broadcast: AI is thinking
      broadcast(room, { type: "ai-start", userId: msg.userId });

      // Stream LLM response
      const result = streamText({
        model: openai("gpt-4o-mini"),
        prompt: msg.prompt,
      });

      for await (const chunk of result.textStream) {
        broadcast(room, {
          type: "ai-chunk",
          content: chunk,
          userId: msg.userId,
        });
      }

      broadcast(room, { type: "ai-end", userId: msg.userId });
    }
  });

  ws.on("close", () => {
    if (currentRoom) rooms.get(currentRoom)?.delete(ws);
  });
});

function broadcast(room: Set<WebSocket>, data: Record<string, unknown>) {
  const msg = JSON.stringify(data);
  room.forEach((client) => {
    if (client.readyState === WebSocket.OPEN) client.send(msg);
  });
}

React Hook: useRealtimeAI

// hooks/use-realtime-ai.ts
"use client";
import { useState, useEffect, useRef, useCallback } from "react";

interface AIMessage {
  type: "ai-start" | "ai-chunk" | "ai-end";
  content?: string;
  userId: string;
}

export function useRealtimeAI(roomId: string) {
  const wsRef = useRef<WebSocket | null>(null);
  const [streaming, setStreaming] = useState(false);
  const [response, setResponse] = useState("");
  const [connected, setConnected] = useState(false);

  useEffect(() => {
    const ws = new WebSocket("ws://localhost:8080");
    wsRef.current = ws;

    ws.onopen = () => {
      setConnected(true);
      ws.send(JSON.stringify({ type: "join", roomId }));
    };

    ws.onmessage = (event) => {
      const msg: AIMessage = JSON.parse(event.data);
      switch (msg.type) {
        case "ai-start":
          setStreaming(true);
          setResponse("");
          break;
        case "ai-chunk":
          setResponse((prev) => prev + (msg.content || ""));
          break;
        case "ai-end":
          setStreaming(false);
          break;
      }
    };

    ws.onclose = () => setConnected(false);

    return () => ws.close();
  }, [roomId]);

  const sendQuery = useCallback((prompt: string, userId: string) => {
    wsRef.current?.send(JSON.stringify({
      type: "ai-query",
      prompt,
      userId,
    }));
  }, []);

  return { connected, streaming, response, sendQuery };
}

Use Cases

• Collaborative AI editor: Nhiều users cùng xem AI generate content real-time
• Live Q&A: AI trả lời questions trong group chat, everyone sees streaming response
• Real-time translation: AI translate messages khi users chat cross-language
• Pair programming: AI assistant stream code suggestions visible cho cả team

Production Considerations

• Reconnection: Implement exponential backoff reconnect khi connection drops
• Authentication: Validate JWT token khi WebSocket connect, reject unauthorized
• Rate limiting: Limit AI queries per user per minute — WebSocket makes this trickier
• Scaling: Single WS server có limit. Dùng Redis pub/sub để sync across multiple instances
• Fallback: Nếu WebSocket unavailable, fallback về Server-Sent Events hoặc polling

Real-time AI features tạo experiences mà traditional request-response không thể match. WebSocket + LLM streaming là foundation cho next generation collaborative AI tools.