feat: show image in user input (#30)

<img width="854" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/789a38c8-f491-456c-ad9a-eb22e19a5bc7">
<img width="869" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/864a6d3e-44c1-4f52-96cd-7f4e55b1178d">
<img width="1288" alt="image"
src="https://github.com/landing-ai/vision-agent-ui/assets/132986242/f8086c87-7a8e-4ec5-bcd5-84572462a964">
https://app.asana.com/0/1204554785675703/1207242625376120/f

app/api/vision-agent/route.ts

@@ -4,6 +4,7 @@ import { StreamingTextResponse } from 'ai';
 import { MessageBase } from '../../../lib/types';
 import { withLogging } from '@/lib/logger';
 import { CLEANED_SEPARATOR } from '@/lib/constants';
+import { cleanAnswerMessage, cleanInputMessage } from '@/lib/messageUtils';
 
 // export const runtime = 'edge';
 export const dynamic = 'force-dynamic';
@@ -33,14 +34,17 @@ export const POST = withLogging(
       JSON.stringify(
         messages.map(message => {
           if (message.role !== 'assistant') {
-            return message;
+            return {
+              ...message,
+              content: cleanInputMessage(message.content),
+            };
           } else {
             const splitedContent = message.content.split(CLEANED_SEPARATOR);
             return {
               ...message,
               content:
                 splitedContent.length > 1
-                  ? splitedContent[1].replace(/!\[answers.*?\.png\)/g, '')
+                  ? cleanAnswerMessage(splitedContent[1])
                   : message.content,
             };
           }

app/chat/page.tsx

@@ -1,6 +1,7 @@
 'use client';
 
 import ImageSelector from '@/components/chat/ImageSelector';
+import { generateInputImageMarkdown } from '@/lib/messageUtils';
 import { ChatEntity, MessageBase } from '@/lib/types';
 import { fetcher } from '@/lib/utils';
 import Image from 'next/image';
@@ -17,7 +18,12 @@ const examples: Example[] = [
     initMessages: [
       {
         role: 'user',
-        content: 'how many cereals are there in the image?',
+        content:
+          'how many cereals are there in the image?' +
+          '\n\n' +
+          generateInputImageMarkdown(
+            'https://landing-lens-support.s3.us-east-2.amazonaws.com/vision-agent-examples/cereal-example.jpg',
+          ),
         id: 'fake-id-1',
       },
     ],

components/chat-sidebar/ChatCard.tsx

@@ -9,6 +9,7 @@ import Image from 'next/image';
 import clsx from 'clsx';
 import Img from '../ui/Img';
 import { format } from 'date-fns';
+import { cleanInputMessage } from '@/lib/messageUtils';
 // import { format } from 'date-fns';
 
 type ChatCardProps = PropsWithChildren<{
@@ -35,7 +36,7 @@ const ChatCard: React.FC<ChatCardProps> = ({ chat }) => {
   const { id: chatIdFromParam } = useParams();
   const pathname = usePathname();
   const { id, url, messages, user, updatedAt } = chat;
-  const firstMessage = messages?.[0]?.content;
+  const firstMessage = cleanInputMessage(messages?.[0]?.content ?? '');
   const title = firstMessage
     ? firstMessage.length > 50
       ? firstMessage.slice(0, 50) + '...'

components/chat/ChatMessage.tsx

@@ -4,26 +4,31 @@
 import remarkGfm from 'remark-gfm';
 import remarkMath from 'remark-math';
 
+import { useMemo } from 'react';
 import { cn } from '@/lib/utils';
 import { CodeBlock } from '@/components/ui/CodeBlock';
 import { MemoizedReactMarkdown } from '@/components/chat/MemoizedReactMarkdown';
 import { IconOpenAI, IconUser } from '@/components/ui/Icons';
-import { ChatMessageActions } from '@/components/chat/ChatMessageActions';
 import { MessageBase } from '../../lib/types';
-import { useCleanedUpMessages } from '@/lib/hooks/useCleanedUpMessages';
 import {
   Tooltip,
   TooltipContent,
   TooltipTrigger,
 } from '@/components/ui/Tooltip';
 import Img from '../ui/Img';
+import { getCleanedUpMessages } from '@/lib/messageUtils';
 
 export interface ChatMessageProps {
   message: MessageBase;
 }
 
 export function ChatMessage({ message, ...props }: ChatMessageProps) {
-  const { logs, content } = useCleanedUpMessages(message);
+  const { logs, content } = useMemo(() => {
+    return getCleanedUpMessages({
+      content: message.content,
+      role: message.role,
+    });
+  }, [message.content, message.role]);
   return (
     <div className={cn('group relative mb-4 flex items-start')} {...props}>
       <div

components/chat/Composer.tsx

@@ -20,6 +20,7 @@ import {
   IconStop,
 } from '@/components/ui/Icons';
 import { cn } from '@/lib/utils';
+import { generateInputImageMarkdown } from '@/lib/messageUtils';
 
 export interface ComposerProps
   extends Pick<
@@ -69,7 +70,8 @@ export function Composer({
             setInput('');
             await append({
               id,
-              content: input,
+              content:
+                input + (url ? '\n\n' + generateInputImageMarkdown(url) : ''),
               role: 'user',
             });
             scrollToBottom();

lib/hooks/{useVisionAgent.tsx → useVisionAgent.ts}

@@ -4,7 +4,11 @@ import { useEffect, useState } from 'react';
 import { ChatEntity, MessageBase, SignedPayload } from '../types';
 import { saveKVChatMessage } from '../kv/chat';
 import { fetcher, nanoid } from '../utils';
-import { getCleanedUpMessages } from './useCleanedUpMessages';
+import {
+  getCleanedUpMessages,
+  generateAnswersImageMarkdown,
+  generateInputImageMarkdown,
+} from '../messageUtils';
 import { CLEANED_SEPARATOR } from '../constants';
 
 const uploadBase64 = async (
@@ -74,8 +78,8 @@ const useVisionAgent = (chat: ChatEntity) => {
         );
         const newContent = publicUrls.reduce((accum, url, index) => {
           return accum.replace(
-            `![answers-${index}](/loading.gif)`,
-            `![answers-${index}](${url})`,
+            generateAnswersImageMarkdown(index, '/loading.gif'),
+            generateAnswersImageMarkdown(index, url),
           );
         }, content);
         const newMessage = {
@@ -93,7 +97,7 @@ const useVisionAgent = (chat: ChatEntity) => {
                 {
                   id: nanoid(),
                   role: 'user',
-                  content: input,
+                  content: input + '\n\n' + generateInputImageMarkdown(url),
                   createdAt: new Date(),
                 } satisfies Message,
               ]

lib/{hooks/useCleanedUpMessages.ts → messageUtils.ts}

@@ -1,6 +1,5 @@
-import { useMemo } from 'react';
-import { MessageBase } from '../types';
-import { CLEANED_SEPARATOR } from '../constants';
+import { MessageBase } from './types';
+import { CLEANED_SEPARATOR } from './constants';
 
 const PAIRS: Record<string, string> = {
   '┍': '┑',
@@ -12,6 +11,26 @@ const PAIRS: Record<string, string> = {
 const MIDDLE_STARTER = '┝';
 const MIDDLE_SEPARATOR = '┿';
 
+const ANSWERS_PREFIX = 'answers';
+const INPUT_PREFIX = 'input';
+
+export const generateAnswersImageMarkdown = (index: number, url: string) => {
+  return `![${ANSWERS_PREFIX}-${index}](${url})`;
+};
+
+export const generateInputImageMarkdown = (url: string, index = 0) => {
+  const prefix = 'input';
+  return `![${INPUT_PREFIX}-${index}](${url})`;
+};
+
+export const cleanInputMessage = (content: string) => {
+  return content.replace(/!\[input-.*?\)/g, '');
+};
+
+export const cleanAnswerMessage = (content: string) => {
+  return content.replace(/!\[answers.*?\.png\)/g, '');
+};
+
 export const getCleanedUpMessages = ({
   content,
   role,
@@ -70,15 +89,11 @@ export const getCleanedUpMessages = ({
     logs: cleanedLogs.join('').replace(/│/g, '|').split('|\n\n|').join('|\n|'),
     content:
       answerText.replace('</</ANSWER>', '').replace('</ANSWER>', '') +
-      images.map((_, index) => `![answers-${index}](/loading.gif)`).join('') +
+      '\n\n' +
+      images
+        .map((_, index) => generateAnswersImageMarkdown(index, '/loading.gif'))
+        .join('') +
       rest.join(''),
     images: images,
   };
 };
-
-export const useCleanedUpMessages = ({ content, role }: MessageBase) => {
-  const cleanedMessage = useMemo(() => {
-    return getCleanedUpMessages({ content, role });
-  }, [content, role]);
-  return cleanedMessage;
-};