lib/hooks/useVisionAgent.ts

import { useChat } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useRef } from 'react';
import { ChatWithMessages, MessageUI } from '../types';
import { convertDBMessageToAPIMessage } from '../utils/message';
import { useSetAtom } from 'jotai';
import { selectedMessageId } from '@/state/chat';
import { Message } from '@prisma/client';
import { useRouter } from 'next/navigation';

const useVisionAgent = (chat: ChatWithMessages) => {
  const { messages: dbMessages, id, mediaUrl } = chat;
  const latestDbMessage = dbMessages[dbMessages.length - 1];
  const setMessageId = useSetAtom(selectedMessageId);

  // Temporary solution for now while single we have to pass mediaUrl separately outside of the messages
  const currMediaUrl = useRef<string>(mediaUrl);
  const currMessageId = useRef<string>(latestDbMessage?.id);
  const router = useRouter();

  const { append, isLoading, data, reload } = useChat({
    api: '/api/vision-agent',
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
    onFinish: message => {
      router.refresh();
      setMessageId(currMessageId.current);
    },
    body: {
      mediaUrl: currMediaUrl.current,
      chatId: id,
      messageId: currMessageId.current,
      // for some reason, the messages has to be stringified to be sent to the API
      apiMessages: JSON.stringify(convertDBMessageToAPIMessage(dbMessages)),
    },
    onError: err => {
      err && toast.error(err.message);
    },
    initialMessages: convertDBMessageToAPIMessage(dbMessages),
  });

  /**
   * If case this is first time user navigated with init message, we need to reload the chat for the first response
   */
  const once = useRef(true);
  useEffect(() => {
    if (!isLoading && !latestDbMessage.result && once.current) {
      once.current = false;
      reload();
    }
  }, [isLoading, latestDbMessage.result, reload]);

  return {
    append: (message: Message) => {
      currMediaUrl.current = message.mediaUrl;
      currMessageId.current = message.id;
      append({
        id,
        role: 'user',
        content: message.prompt,
      });
    },
    data: data as unknown as PrismaJson.MessageBody[],
    reload,
    isLoading,
  };
};

export default useVisionAgent;