chat-proxy/frontend/src/useChatStream.js

// useChatStream.js
import { useState, useCallback, useRef } from "react";
import { getSessionId, getUserId } from './useSessionId';

export function useChatStream() {
  const [messages, setMessages] = useState([]);
  const [loading, setLoading] = useState(false);
  const abortRef = useRef(null);
  const sessionId = getSessionId();
  const userId = getUserId();

  const sendMessage = useCallback(
    async (input) => {
      if (!input.trim()) return;

      // Abort any previous stream
      if (abortRef.current) {
        abortRef.current.abort();
      }

      const controller = new AbortController();
      abortRef.current = controller;

      const userMessage = { role: "user", content: input };
      const assistantIndex = messages.length + 1;
      setMessages((prev) => [...prev, userMessage, { role: "assistant", content: "" }]);
      setLoading(true);

      try {
        const res = await fetch("/v1/chat-stream", {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            Accept: "text/event-stream",
          },
	  body: JSON.stringify({ user_id: userId, session_id: sessionId, message: input }),
          signal: controller.signal,
        });

        if (!res.ok || !res.body) throw new Error(`HTTP ${res.status}`);

        const reader = res.body.getReader();
        const decoder = new TextDecoder("utf-8");
        let buffer = "";
        let acc = "";

        while (true) {
          const { value, done } = await reader.read();
          if (done) break;
          buffer += decoder.decode(value, { stream: true });
          const parts = buffer.split("\n\n");
          buffer = parts.pop() || "";

          for (const part of parts) {
            const line = part
              .split("\n")
              .map((l) => l.trim())
              .find((l) => l.startsWith("data:"));
            if (!line) continue;

            const data = line.slice(5).trim();
            if (data === "[DONE]") {
              setMessages((prev) => {
                const next = [...prev];
                next[assistantIndex] = { role: "assistant", content: acc };
                return next;
              });
              setLoading(false);
              return;
            }

            try {
              const obj = JSON.parse(data);
              const piece = obj?.choices?.[0]?.delta?.content ?? obj?.choices?.[0]?.text ?? "";
              if (piece) {
                acc += piece;
                setMessages((prev) => {
                  const next = [...prev];
                  next[assistantIndex] = { role: "assistant", content: acc };
                  return next;
                });
              }
            } catch {
              /* ignore malformed JSON chunks */
            }
          }
        }
      } catch (err) {
        if (err.name !== "AbortError") {
          setMessages((prev) => {
            const next = [...prev];
            next[assistantIndex] = {
              role: "assistant",
              content: `Error: ${String(err)}`,
            };
            return next;
          });
        }
      } finally {
        setLoading(false);
        abortRef.current = null;
      }
    },
    [messages]
  );

  const stopGenerating = useCallback(() => {
    if (abortRef.current) {
      abortRef.current.abort();
      abortRef.current = null;
      setLoading(false);
    }
  }, []);

  return { messages, loading, sendMessage, stopGenerating, setMessages };
}