lib/hooks/useVisionAgent.ts

import { useChat, type Message, UseChatHelpers } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useState } from 'react';
import { MessageBase, SignedPayload } from '../types';
import { fetcher, nanoid } from '../utils';
import {
  getCleanedUpMessages,
  generateAnswersImageMarkdown,
  generateInputImageMarkdown,
} from '../messageUtils';
import { CLEANED_SEPARATOR } from '../constants';
import { useSearchParams } from 'next/navigation';
import { ChatWithMessages, MessageRaw } from '../db/types';
import { dbPostCreateMessage } from '../db/functions';

const uploadBase64 = async (
  base64: string,
  messageId: string,
  chatId: string,
  index: number,
) => {
  const res = await fetch(
    'data:image/png;base64,' + base64.replace('base:64', ''),
  );
  const blob = await res.blob();
  const { signedUrl, publicUrl, fields } = await fetcher<SignedPayload>(
    '/api/sign',
    {
      method: 'POST',
      body: JSON.stringify({
        id: `${chatId}/${messageId}`,
        fileType: blob.type,
        fileName: `answer-${index}.${blob.type.split('/')[1]}`,
      }),
    },
  );
  const formData = new FormData();
  Object.entries(fields).forEach(([key, value]) => {
    formData.append(key, value as string);
  });
  formData.append('file', blob);

  const uploadResponse = await fetch(signedUrl, {
    method: 'POST',
    body: formData,
  });
  if (uploadResponse.ok) {
    return publicUrl;
  } else {
    throw new Error('Upload failed');
  }
};

const useVisionAgent = (chat: ChatWithMessages) => {
  const { messages: initialMessages, id, mediaUrl } = chat;
  const searchParams = useSearchParams();
  const reflectionValue = searchParams.get('reflection');

  const {
    messages,
    append: appendRaw,
    reload,
    stop,
    isLoading,
    input,
    setInput,
    setMessages,
    error,
  } = useChat({
    api: '/api/vision-agent',
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
    onFinish: async message => {
      const { logs = '', content, images } = getCleanedUpMessages(message);
      if (images?.length) {
        const publicUrls = await Promise.all(
          images.map((image, index) =>
            uploadBase64(image, message.id, id ?? 'no-id', index),
          ),
        );
        const newContent = publicUrls.reduce((accum, url, index) => {
          return accum.replace(
            generateAnswersImageMarkdown(index, '/loading.gif'),
            generateAnswersImageMarkdown(index, url),
          );
        }, content);
        const newMessage = {
          ...message,
          content: logs + CLEANED_SEPARATOR + newContent,
        };
        setMessages([
          ...messages,
          /**
           * A workaround to fix the issue of the messages been stale state when appending a new message
           * https://github.com/vercel/ai/issues/550#issuecomment-1712693371
           */
          ...(input
            ? [
                {
                  id: nanoid(),
                  role: 'user',
                  content:
                    input + '\n\n' + generateInputImageMarkdown(mediaUrl),
                  createdAt: new Date(),
                } satisfies Message,
              ]
            : []),
          newMessage,
        ]);
        await dbPostCreateMessage(id, {
          role: newMessage.role as 'user' | 'assistant',
          content: newMessage.content,
        });
      } else {
        await dbPostCreateMessage(id, {
          role: message.role as 'user' | 'assistant',
          content: logs + CLEANED_SEPARATOR + content,
        });
      }
    },
    initialMessages: initialMessages,
    body: {
      mediaUrl,
      id,
      enableSelfReflection: reflectionValue === 'true',
    },
  });

  /**
   * If the last message is from the user, reload the chat, this would trigger to get the response from the assistant
   * There are 2 scenarios when this might happen
   * 1. Navigated from example images, init message only include preset user message
   * 2. Last time the assistant message failed or not saved to database.
   */
  useEffect(() => {
    if (
      !isLoading &&
      messages.length &&
      messages[messages.length - 1].role === 'user'
    ) {
      reload();
    }
  }, [isLoading, messages, reload]);

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

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

export default useVisionAgent;