lib/hooks/useVisionAgent.ts

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

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

  // Temporary solution for now while single we have to pass mediaUrl separately outside of the messages
  const router = useRouter();
  const setMessageId = useSetAtom(selectedMessageId);

  const {
    data = [],
    reload,
    append,
    messages,
    isLoading,
  } = useChat({
    api: '/api/vision-agent',
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
    onFinish: () => {
      setMessageId(latestDbMessage.id);
      router.refresh();
    },
    body: {
      mediaUrl: latestDbMessage.mediaUrl,
      chatId: id,
      messageId: latestDbMessage.id,
      // 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);
    },
  });

  /**
   * 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(() => {
    const appendDbMessage = async (latestDbMessage: Message) => {
      await append({
        id: latestDbMessage.id + '-user',
        content: latestDbMessage.prompt,
        role: 'user',
      });
    };
    if (isLoading || latestDbMessage.response || latestDbMessage.responseBody) {
      return;
    }
    if (messages.length === 0) {
      if (once.current) {
        once.current = false;
        appendDbMessage(latestDbMessage);
      }
      return;
    }
    if (
      messages.findIndex(message => message.id.includes(latestDbMessage.id)) ===
      -1
    ) {
      appendDbMessage(latestDbMessage);
    }
  }, [latestDbMessage, messages, isLoading]);

  const initDataIndex = data.findIndex(
    (m: any) =>
      m.type === 'init' && m.payload?.messageId === latestDbMessage.id,
  );

  return {
    data:
      initDataIndex >= 0
        ? (data.slice(initDataIndex + 1) as unknown as PrismaJson.MessageBody[])
        : [],
    reload,
    isLoading,
  };
};

export default useVisionAgent;