lib/hooks/useChatWithDataset.ts

import { useChat, type Message } from 'ai/react';
import { useAtom } from 'jotai';
import { toast } from 'react-hot-toast';
import { datasetAtom } from '../../state';
import { useEffect, useState } from 'react';
import { MessageWithSelectedDataset } from '../types';

const useChatWithDataset = () => {
  const [dataset] = useAtom(datasetAtom);
  const {
    messages,
    append,
    reload,
    stop,
    isLoading,
    input,
    setInput,
    setMessages,
  } = useChat({
    sendExtraMessageFields: true,
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
  });

  const [loadingDots, setLoadingDots] = useState('');

  useEffect(() => {
    let loadingInterval: NodeJS.Timeout;

    if (isLoading) {
      loadingInterval = setInterval(() => {
        setLoadingDots(prevMessage => {
          switch (prevMessage) {
            case '':
              return '.';
            case '.':
              return '..';
            case '..':
              return '...';
            case '...':
              return '';
            default:
              return '';
          }
        });
      }, 500);
    }

    return () => {
      clearInterval(loadingInterval);
    };
  }, [isLoading]);

  const assistantLoadingMessage = {
    id: 'loading',
    content: loadingDots,
    role: 'assistant',
  };

  const messageWithLoading =
    isLoading &&
    messages.length &&
    messages[messages.length - 1].role !== 'assistant'
      ? [...messages, assistantLoadingMessage]
      : messages;

  const selectedDataset = dataset.find(entity => entity.selected)
    ? dataset.filter(entity => entity.selected)
    : // If there is no selected dataset, use the entire dataset
      dataset;

  const appendWithDataset: typeof append = message => {
    // const newSystemMessage: Message = {
    // 	id: 'fake-id',
    // 	content:
    // 		'For the next prompt, here are names of images provided by user, please use these name if you need reference: ' +
    // 		selectedDataset.map(entity => entity.name).join(', '),
    // 	role: 'system',
    // };
    // const newSystemMessage: Message = {
    // 	id: 'fake-id',
    // 	content: `For the next prompt, please use tags provided by the user to assign to corresponding images.
    // 		For example:

    // 		Input:
    // 		red, blue, round

    // 		Answer (each in a new line):
    // 		Image 1: red\n
    // 		Image 2: blue,round\n`,
    // 	role: 'system',
    // };
    // setMessages([...messages, newSystemMessage]);
    return append({
      ...message,
      // @ts-ignore this is extra fields
      dataset: selectedDataset,
    } as MessageWithSelectedDataset);
  };

  return {
    messages: messageWithLoading as MessageWithSelectedDataset[],
    append: appendWithDataset,
    reload,
    stop,
    isLoading,
    input,
    setInput,
  };
};

export default useChatWithDataset;