multi image

app/api/chat/route.ts

@@ -4,11 +4,11 @@ import OpenAI from 'openai';
 import { auth } from '@/auth';
 import { nanoid } from '@/lib/utils';
 import {
+	ChatCompletionMessageParam,
 	ChatCompletionContentPart,
 	ChatCompletionContentPartImage,
-	ChatCompletionMessageParam,
 } from 'openai/resources';
-import { DatasetImageEntity } from '../../../lib/types';
+import { MessageWithSelectedDataset } from '../../../lib/types';
 
 export const runtime = 'edge';
 
@@ -18,9 +18,8 @@ const openai = new OpenAI({
 
 export async function POST(req: Request) {
 	const json = await req.json();
-	const { messages, dataset } = json as {
-		messages: ChatCompletionMessageParam[];
-		dataset: DatasetImageEntity[];
+	const { messages } = json as {
+		messages: MessageWithSelectedDataset[];
 	};
 
 	const userId = (await auth())?.user.id;
@@ -31,15 +30,16 @@ export async function POST(req: Request) {
 		});
 	}
 
-	const messagesWithImage: ChatCompletionMessageParam[] = messages.map(
+	const formattedMessage: ChatCompletionMessageParam[] = messages.map(
 		message => {
-			if (message.role !== 'user') return message;
+			const { dataset, ...rest } = message;
+
 			const contentWithImage: ChatCompletionContentPart[] = [
 				{
 					type: 'text',
 					text: message.content as string,
 				},
-				...dataset.map(
+				...(dataset ?? []).map(
 					entity =>
 						({
 							type: 'image_url',
@@ -56,32 +56,13 @@ export async function POST(req: Request) {
 
 	const res = await openai.chat.completions.create({
 		model: 'gpt-4-vision-preview',
-		messages: messagesWithImage,
-		temperature: 0.7,
+		messages: formattedMessage,
+		temperature: 0.3,
 		stream: true,
 		max_tokens: 300,
 	});
 
-	const stream = OpenAIStream(res, {
-		async onCompletion(completion) {
-			const title = json.messages[0].content.substring(0, 100);
-			const id = json.id ?? nanoid();
-			const createdAt = Date.now();
-			const payload = {
-				id,
-				title,
-				userId,
-				createdAt,
-				messages: [
-					...messages,
-					{
-						content: completion,
-						role: 'assistant',
-					},
-				],
-			};
-		},
-	});
+	const stream = OpenAIStream(res);
 
 	return new StreamingTextResponse(stream);
 }

app/globals.css

@@ -87,8 +87,8 @@
 		height: 50px;
 		background: linear-gradient(
 			to bottom,
-			rgba(255, 255, 255, 0),
-			rgba(255, 255, 255, 1)
+			rgba(255, 255, 255, 1),
+			rgba(255, 255, 255, 0)
 		);
 		pointer-events: none;
 	}

components/chat-message-actions.tsx

@@ -1,40 +1,41 @@
-'use client'
+'use client';
 
-import { type Message } from 'ai'
+import { type Message } from 'ai';
 
-import { Button } from '@/components/ui/button'
-import { IconCheck, IconCopy } from '@/components/ui/icons'
-import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard'
-import { cn } from '@/lib/utils'
+import { Button } from '@/components/ui/button';
+import { IconCheck, IconCopy } from '@/components/ui/icons';
+import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard';
+import { cn } from '@/lib/utils';
+import { MessageWithSelectedDataset } from '../lib/types';
 
 interface ChatMessageActionsProps extends React.ComponentProps<'div'> {
-  message: Message
+	message: MessageWithSelectedDataset;
 }
 
 export function ChatMessageActions({
-  message,
-  className,
-  ...props
+	message,
+	className,
+	...props
 }: ChatMessageActionsProps) {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 })
+	const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 });
 
-  const onCopy = () => {
-    if (isCopied) return
-    copyToClipboard(message.content)
-  }
+	const onCopy = () => {
+		if (isCopied) return;
+		copyToClipboard(message.content);
+	};
 
-  return (
-    <div
-      className={cn(
-        'flex items-center justify-end transition-opacity group-hover:opacity-100 md:absolute md:-right-10 md:-top-2 md:opacity-0',
-        className
-      )}
-      {...props}
-    >
-      <Button variant="ghost" size="icon" onClick={onCopy}>
-        {isCopied ? <IconCheck /> : <IconCopy />}
-        <span className="sr-only">Copy message</span>
-      </Button>
-    </div>
-  )
+	return (
+		<div
+			className={cn(
+				'flex items-center justify-end transition-opacity group-hover:opacity-100 md:absolute md:-right-10 md:-top-2 md:opacity-0',
+				className,
+			)}
+			{...props}
+		>
+			<Button variant="ghost" size="icon" onClick={onCopy}>
+				{isCopied ? <IconCheck /> : <IconCopy />}
+				<span className="sr-only">Copy message</span>
+			</Button>
+		</div>
+	);
 }

components/chat-message.tsx

@@ -1,7 +1,6 @@
 // Inspired by Chatbot-UI and modified to fit the needs of this project
 // @see https://github.com/mckaywrigley/chatbot-ui/blob/main/components/Chat/ChatMessage.tsx
 
-import { Message } from 'ai';
 import remarkGfm from 'remark-gfm';
 import remarkMath from 'remark-math';
 
@@ -10,9 +9,10 @@ import { CodeBlock } from '@/components/ui/codeblock';
 import { MemoizedReactMarkdown } from '@/components/markdown';
 import { IconOpenAI, IconUser } from '@/components/ui/icons';
 import { ChatMessageActions } from '@/components/chat-message-actions';
+import { MessageWithSelectedDataset } from '../lib/types';
 
 export interface ChatMessageProps {
-	message: Message;
+	message: MessageWithSelectedDataset;
 }
 
 export function ChatMessage({ message, ...props }: ChatMessageProps) {

components/chat-panel.tsx

@@ -7,20 +7,16 @@ import { PromptForm } from '@/components/prompt-form';
 import { ButtonScrollToBottom } from '@/components/button-scroll-to-bottom';
 import { IconRefresh, IconShare, IconStop } from '@/components/ui/icons';
 import { ChatShareDialog } from '@/components/chat-share-dialog';
+import { MessageWithSelectedDataset } from '../lib/types';
 
 export interface ChatPanelProps
 	extends Pick<
 		UseChatHelpers,
-		| 'append'
-		| 'isLoading'
-		| 'reload'
-		| 'messages'
-		| 'stop'
-		| 'input'
-		| 'setInput'
+		'append' | 'isLoading' | 'reload' | 'stop' | 'input' | 'setInput'
 	> {
 	id?: string;
 	title?: string;
+	messages: MessageWithSelectedDataset[];
 }
 
 export function ChatPanel({
@@ -37,7 +33,7 @@ export function ChatPanel({
 	const [shareDialogOpen, setShareDialogOpen] = React.useState(false);
 
 	return (
-		<div className="fixed inset-x-0 bottom-0 w-full bg-gradient-to-b from-muted/30 from-0% to-muted/30 to-50% animate-in duration-300 ease-in-out dark:from-background/10 dark:from-10% dark:to-background/80 peer-[[data-state=open]]:group-[]:lg:pl-[250px] peer-[[data-state=open]]:group-[]:xl:pl-[300px]">
+		<div className="fixed inset-x-0 bottom-0 w-full animate-in duration-300 ease-in-out peer-[[data-state=open]]:group-[]:lg:pl-[250px] peer-[[data-state=open]]:group-[]:xl:pl-[300px]">
 			<ButtonScrollToBottom />
 			<div className="mx-auto sm:max-w-3xl sm:px-4">
 				<div className="flex items-center justify-center h-12">
@@ -57,7 +53,7 @@ export function ChatPanel({
 									<IconRefresh className="mr-2" />
 									Regenerate response
 								</Button>
-								{id && title ? (
+								{/* {id && title ? (
 									<>
 										<Button
 											variant="outline"
@@ -78,7 +74,7 @@ export function ChatPanel({
 											}}
 										/>
 									</>
-								) : null}
+								) : null} */}
 							</div>
 						)
 					)}

components/chat-share-dialog.tsx

@@ -1,106 +1,106 @@
-'use client'
+'use client';
 
-import * as React from 'react'
-import { type DialogProps } from '@radix-ui/react-dialog'
-import { toast } from 'react-hot-toast'
+import * as React from 'react';
+import { type DialogProps } from '@radix-ui/react-dialog';
+import { toast } from 'react-hot-toast';
 
-import { ServerActionResult, type Chat } from '@/lib/types'
-import { Button } from '@/components/ui/button'
+import { ServerActionResult, type Chat } from '@/lib/types';
+import { Button } from '@/components/ui/button';
 import {
-  Dialog,
-  DialogContent,
-  DialogDescription,
-  DialogFooter,
-  DialogHeader,
-  DialogTitle
-} from '@/components/ui/dialog'
-import { IconSpinner } from '@/components/ui/icons'
-import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard'
+	Dialog,
+	DialogContent,
+	DialogDescription,
+	DialogFooter,
+	DialogHeader,
+	DialogTitle,
+} from '@/components/ui/dialog';
+import { IconSpinner } from '@/components/ui/icons';
+import { useCopyToClipboard } from '@/lib/hooks/use-copy-to-clipboard';
 
 interface ChatShareDialogProps extends DialogProps {
-  chat: Pick<Chat, 'id' | 'title' | 'messages'>
-  shareChat: (id: string) => ServerActionResult<Chat>
-  onCopy: () => void
+	chat: Pick<Chat, 'id' | 'title' | 'messages'>;
+	shareChat: (id: string) => ServerActionResult<Chat>;
+	onCopy: () => void;
 }
 
 export function ChatShareDialog({
-  chat,
-  shareChat,
-  onCopy,
-  ...props
+	chat,
+	shareChat,
+	onCopy,
+	...props
 }: ChatShareDialogProps) {
-  const { copyToClipboard } = useCopyToClipboard({ timeout: 1000 })
-  const [isSharePending, startShareTransition] = React.useTransition()
+	const { copyToClipboard } = useCopyToClipboard({ timeout: 1000 });
+	const [isSharePending, startShareTransition] = React.useTransition();
 
-  const copyShareLink = React.useCallback(
-    async (chat: Chat) => {
-      if (!chat.sharePath) {
-        return toast.error('Could not copy share link to clipboard')
-      }
+	const copyShareLink = React.useCallback(
+		async (chat: Chat) => {
+			if (!chat.sharePath) {
+				return toast.error('Could not copy share link to clipboard');
+			}
 
-      const url = new URL(window.location.href)
-      url.pathname = chat.sharePath
-      copyToClipboard(url.toString())
-      onCopy()
-      toast.success('Share link copied to clipboard', {
-        style: {
-          borderRadius: '10px',
-          background: '#333',
-          color: '#fff',
-          fontSize: '14px'
-        },
-        iconTheme: {
-          primary: 'white',
-          secondary: 'black'
-        }
-      })
-    },
-    [copyToClipboard, onCopy]
-  )
+			const url = new URL(window.location.href);
+			url.pathname = chat.sharePath;
+			copyToClipboard(url.toString());
+			onCopy();
+			toast.success('Share link copied to clipboard', {
+				style: {
+					borderRadius: '10px',
+					background: '#333',
+					color: '#fff',
+					fontSize: '14px',
+				},
+				iconTheme: {
+					primary: 'white',
+					secondary: 'black',
+				},
+			});
+		},
+		[copyToClipboard, onCopy],
+	);
 
-  return (
-    <Dialog {...props}>
-      <DialogContent>
-        <DialogHeader>
-          <DialogTitle>Share link to chat</DialogTitle>
-          <DialogDescription>
-            Anyone with the URL will be able to view the shared chat.
-          </DialogDescription>
-        </DialogHeader>
-        <div className="p-4 space-y-1 text-sm border rounded-md">
-          <div className="font-medium">{chat.title}</div>
-          <div className="text-muted-foreground">
-            {chat.messages.length} messages
-          </div>
-        </div>
-        <DialogFooter className="items-center">
-          <Button
-            disabled={isSharePending}
-            onClick={() => {
-              // @ts-ignore
-              startShareTransition(async () => {
-                const result = await shareChat(chat.id)
+	return (
+		<Dialog {...props}>
+			<DialogContent>
+				<DialogHeader>
+					<DialogTitle>Share link to chat</DialogTitle>
+					<DialogDescription>
+						Anyone with the URL will be able to view the shared chat.
+					</DialogDescription>
+				</DialogHeader>
+				<div className="p-4 space-y-1 text-sm border rounded-md">
+					<div className="font-medium">{chat.title}</div>
+					<div className="text-muted-foreground">
+						{chat.messages.length} messages
+					</div>
+				</div>
+				<DialogFooter className="items-center">
+					<Button
+						disabled={isSharePending}
+						onClick={() => {
+							// @ts-ignore
+							startShareTransition(async () => {
+								const result = await shareChat(chat.id);
 
-                if (result && 'error' in result) {
-                  toast.error(result.error)
-                  return
-                }
+								if (result && 'error' in result) {
+									toast.error(result.error);
+									return;
+								}
 
-                copyShareLink(result)
-              })
-            }}
-          >
-            {isSharePending ? (
-              <>
-                <IconSpinner className="mr-2 animate-spin" />
-                Copying...
-              </>
-            ) : (
-              <>Copy link</>
-            )}
-          </Button>
-        </DialogFooter>
-      </DialogContent>
-    </Dialog>
-  )
+								copyShareLink(result);
+							});
+						}}
+					>
+						{isSharePending ? (
+							<>
+								<IconSpinner className="mr-2 animate-spin" />
+								Copying...
+							</>
+						) : (
+							<>Copy link</>
+						)}
+					</Button>
+				</DialogFooter>
+			</DialogContent>
+		</Dialog>
+	);
 }

components/chat/ChatList.tsx

@@ -1,71 +1,26 @@
 'use client';
 
-import { type Message } from 'ai';
-import { useEffect, useState } from 'react';
-
 import { Separator } from '@/components/ui/separator';
 import { ChatMessage } from '@/components/chat-message';
+import { MessageWithSelectedDataset } from '../../lib/types';
 
 export interface ChatList {
-	messages: Message[];
-	isLoading: boolean;
+	messages: MessageWithSelectedDataset[];
 }
 
-export function ChatList({ messages, isLoading }: ChatList) {
-	const [loadingDots, setLoadingDots] = useState('');
-
-	useEffect(() => {
-		let loadingInterval: NodeJS.Timeout;
-
-		if (isLoading) {
-			loadingInterval = setInterval(() => {
-				setLoadingDots(prevMessage => {
-					switch (prevMessage) {
-						case '':
-							return '.';
-						case '.':
-							return '..';
-						case '..':
-							return '...';
-						case '...':
-							return '';
-						default:
-							return '';
-					}
-				});
-			}, 500);
-		}
-
-		return () => {
-			clearInterval(loadingInterval);
-		};
-	}, [isLoading]);
-
-	if (!messages.length) {
-		return null;
-	}
-
-	const assistantLoadingMessage: Message = {
-		id: 'loading',
-		content: loadingDots,
-		role: 'assistant',
-	};
-
-	const messageWithLoading =
-		isLoading && messages[messages.length - 1].role !== 'assistant'
-			? [...messages, assistantLoadingMessage]
-			: messages;
-
+export function ChatList({ messages }: ChatList) {
 	return (
 		<div className="relative mx-auto max-w-3xl px-8 pr-12">
-			{messageWithLoading.map((message, index) => (
-				<div key={index}>
-					<ChatMessage message={message} />
-					{index < messageWithLoading.length - 1 && (
-						<Separator className="my-4 md:my-8" />
-					)}
-				</div>
-			))}
+			{messages
+				.filter(message => message.role !== 'system')
+				.map((message, index) => (
+					<div key={index}>
+						<ChatMessage message={message} />
+						{index < messages.length - 1 && (
+							<Separator className="my-4 md:my-8" />
+						)}
+					</div>
+				))}
 		</div>
 	);
 }

components/chat/index.tsx

@@ -1,47 +1,28 @@
 'use client';
-
-import { useChat, type Message } from 'ai/react';
-
 import { cn } from '@/lib/utils';
 import { ChatList } from '@/components/chat/ChatList';
 import { ChatPanel } from '@/components/chat-panel';
 import { ChatScrollAnchor } from '@/components/chat-scroll-anchor';
-import { toast } from 'react-hot-toast';
-import { useAtom } from 'jotai';
-import { datasetAtom } from '@/state';
 import ImageList from './ImageList';
+import useChatWithDataset from '../../lib/hooks/useChatWithDataset';
 
 export interface ChatProps extends React.ComponentProps<'div'> {
-	initialMessages?: Message[];
 	id?: string;
 }
 
-export function Chat({ id, initialMessages, className }: ChatProps) {
-	const [dataset] = useAtom(datasetAtom);
+export function Chat({ id, className }: ChatProps) {
 	const { messages, append, reload, stop, isLoading, input, setInput } =
-		useChat({
-			initialMessages,
-			id,
-			body: {
-				id,
-				dataset: dataset,
-			},
-			onResponse(response) {
-				if (response.status === 401) {
-					toast.error(response.statusText);
-				}
-			},
-		});
+		useChatWithDataset();
 
 	return (
 		<>
 			<div className={cn('pb-[150px] pt-4 md:pt-10 h-full', className)}>
 				<div className="flex h-full">
-					<div className="w-1/2 relative border-r-2 border-gray-200">
+					<div className="w-1/2 relative border-r-2 border-gray-300 overflow-auto">
 						<ImageList />
 					</div>
 					<div className="w-1/2 relative overflow-auto">
-						<ChatList messages={messages} isLoading={isLoading} />
+						<ChatList messages={messages} />
 						<ChatScrollAnchor trackVisibility={isLoading} />
 					</div>
 				</div>

lib/hooks/useChatWithDataset.ts

@@ -0,0 +1,102 @@
+import { useChat, type Message } from 'ai/react';
+import { useAtom } from 'jotai';
+import { toast } from 'react-hot-toast';
+import { datasetAtom } from '../../state';
+import { useEffect, useState } from 'react';
+import { MessageWithSelectedDataset } from '../types';
+
+const useChatWithDataset = () => {
+	const [dataset] = useAtom(datasetAtom);
+	const {
+		messages,
+		append,
+		reload,
+		stop,
+		isLoading,
+		input,
+		setInput,
+		setMessages,
+	} = useChat({
+		sendExtraMessageFields: true,
+		onResponse(response) {
+			if (response.status !== 200) {
+				toast.error(response.statusText);
+			}
+		},
+	});
+
+	const [loadingDots, setLoadingDots] = useState('');
+
+	useEffect(() => {
+		let loadingInterval: NodeJS.Timeout;
+
+		if (isLoading) {
+			loadingInterval = setInterval(() => {
+				setLoadingDots(prevMessage => {
+					switch (prevMessage) {
+						case '':
+							return '.';
+						case '.':
+							return '..';
+						case '..':
+							return '...';
+						case '...':
+							return '';
+						default:
+							return '';
+					}
+				});
+			}, 500);
+		}
+
+		return () => {
+			clearInterval(loadingInterval);
+		};
+	}, [isLoading]);
+
+	const assistantLoadingMessage = {
+		id: 'loading',
+		content: loadingDots,
+		role: 'assistant',
+	};
+
+	const messageWithLoading =
+		isLoading &&
+		messages.length &&
+		messages[messages.length - 1].role !== 'assistant'
+			? [...messages, assistantLoadingMessage]
+			: messages;
+
+	const selectedDataset = dataset.find(entity => entity.selected)
+		? dataset.filter(entity => entity.selected)
+		: // If there is no selected dataset, use the entire dataset
+			dataset;
+
+	const appendWithDataset: typeof append = message => {
+		const newSystemMessage: Message = {
+			id: 'fake-id',
+			content:
+				'For the next prompt, here are names of images provided by user, please use these name if you need reference: ' +
+				selectedDataset.map(entity => entity.name).join(', '),
+			role: 'system',
+		};
+		setMessages([...messages, newSystemMessage]);
+		return append({
+			...message,
+			// @ts-ignore this is extra fields
+			dataset: selectedDataset,
+		} satisfies MessageWithSelectedDataset);
+	};
+
+	return {
+		messages: messageWithLoading as MessageWithSelectedDataset[],
+		append: appendWithDataset,
+		reload,
+		stop,
+		isLoading,
+		input,
+		setInput,
+	};
+};
+
+export default useChatWithDataset;

lib/types.ts

@@ -1,4 +1,5 @@
 import { type Message } from 'ai';
+import { CreateMessage } from 'ai/react/dist';
 
 export interface Chat extends Record<string, any> {
 	id: string;
@@ -22,3 +23,7 @@ export type DatasetImageEntity = {
 	selected: boolean;
 	name: string;
 };
+
+export type MessageWithSelectedDataset = (Message | CreateMessage) & {
+	dataset: DatasetImageEntity[];
+};