lib/hooks/useVisionAgent.ts

import { useChat, UseChatHelpers } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useRef } from 'react';
import { ChatWithMessages } from '../db/types';
import { dbPostCreateMessage } from '../db/functions';
import {
  convertDbMessageToMessage,
  convertMessageToDbMessage,
} from '../messageUtils';

const useVisionAgent = (chat: ChatWithMessages) => {
  const { messages: initialMessages, id, mediaUrl } = chat;

  const {
    messages,
    append: appendRaw,
    isLoading,
    reload,
  } = useChat({
    api: '/api/vision-agent',
    // @ts-ignore https://sdk.vercel.ai/docs/troubleshooting/common-issues/use-chat-failed-to-parse-stream
    streamMode: 'text',
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
    onFinish: async message => {
      await dbPostCreateMessage(id, convertMessageToDbMessage(message));
    },
    initialMessages: initialMessages.map(convertDbMessageToMessage),
    body: {
      mediaUrl,
      id,
    },
    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(() => {
    if (
      !isLoading &&
      messages.length === 1 &&
      messages[0].role === 'user' &&
      once.current
    ) {
      once.current = false;
      reload();
    }
  }, [isLoading, messages, reload]);

  const append: UseChatHelpers['append'] = async message => {
    dbPostCreateMessage(id, {
      role: message.role as 'user' | 'assistant',
      content: message.content,
      result: null,
    });
    return appendRaw(message);
  };

  return {
    messages,
    append,
    reload,
    isLoading,
  };
};

export default useVisionAgent;