chatbot-template/app/use-prompt.tsx

import { useState, useCallback, useRef, useEffect } from "react";
import { type Message } from "@prisma/client";
import { nanoid } from "@/lib/utils";

export function usePrompt({
  messages = [],
  id,
}: {
  messages?: Message[];
  id: string | undefined;
}) {
  const [isLoading, setIsLoading] = useState(false);
  const [messageList, setMessageList] = useState(messages);

  const isLoadingRef = useRef(isLoading);
  const messageListRef = useRef(messageList);

  useEffect(() => {
    isLoadingRef.current = isLoading;
  }, [isLoading]);

  useEffect(() => {
    messageListRef.current = messageList;
  }, [messageList]);

  const appendUserMessage = useCallback(
    async (content: string | Message) => {
      // Prevent multiple requests at once
      if (isLoadingRef.current) return;

      const userMsg =
        typeof content === "string"
          ? ({ id: nanoid(10), role: "user", content } as Message)
          : content;
      const assMsg = {
        id: nanoid(10),
        role: "assistant",
        content: "",
      } as Message;
      const messageListSnapshot = messageListRef.current;

      // Reset output
      setIsLoading(true);

      try {
        // Set user input immediately
        setMessageList([...messageListSnapshot, userMsg]);

        // If streaming, we need to use fetchEventSource directly
        const response = await fetch(`/api/generate`, {
          method: "POST",
          body: JSON.stringify({
            id: id || nanoid(10),
            messages: [...messageListSnapshot, userMsg].map((m) => ({
              role: m.role,
              content: m.content,
            })),
          }),
          headers: { "Content-Type": "application/json" },
        });
        // This data is a ReadableStream
        const data = response.body;
        if (!data) {
          return;
        }

        const reader = data.getReader();
        const decoder = new TextDecoder();
        let done = false;
        let accumulatedValue = ""; // Variable to accumulate chunks

        while (!done) {
          const { value, done: doneReading } = await reader.read();
          done = doneReading;
          const chunkValue = decoder.decode(value);
          accumulatedValue += chunkValue; // Accumulate the chunk value

          // Check if the accumulated value contains the delimiter
          const delimiter = "\n";
          const chunks = accumulatedValue.split(delimiter);

          // Process all chunks except the last one (which may be incomplete)
          while (chunks.length > 1) {
            const chunkToDispatch = chunks.shift(); // Get the first chunk
            if (chunkToDispatch && chunkToDispatch.length > 0) {
              const chunk = JSON.parse(chunkToDispatch);
              assMsg.content += chunk;
              setMessageList([...messageListSnapshot, userMsg, assMsg]);
            }
          }

          // The last chunk may be incomplete, so keep it in the accumulated value
          accumulatedValue = chunks[0];
        }

        // Process any remaining accumulated value after the loop is done
        if (accumulatedValue.length > 0) {
          assMsg.content += accumulatedValue;
          setMessageList([...messageListSnapshot, userMsg, assMsg]);
        }
      } finally {
        setIsLoading(false);
      }
    },
    [id]
  );

  const reloadLastMessage = useCallback(async () => {
    // Prevent multiple requests at once
    if (isLoadingRef.current) return;

    const userMsg = messageListRef.current.at(-2);
    const assMsg = messageListRef.current.at(-1);

    // Both should exist.
    if (!userMsg || !assMsg) return;

    messageListRef.current = messageListRef.current.slice(0, -2);
    setMessageList(messageListRef.current);

    await appendUserMessage(userMsg);
  }, [appendUserMessage]);

  return {
    messageList,
    appendUserMessage,
    reloadLastMessage,
    isLoading,
  };
}