lib/hooks/useChatWithMedia.ts

import { useChat, type Message } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useState } from 'react';
import { MediaDetails } from '../fetch';
import { MessageWithSelectedDataset } from '../types';

const useChatWithMedia = (mediaList: MediaDetails[]) => {
	const {
		messages,
		append,
		reload,
		stop,
		isLoading,
		input,
		setInput,
		setMessages,
	} = useChat({
		sendExtraMessageFields: true,
		onResponse(response) {
			if (response.status !== 200) {
				toast.error(response.statusText);
			}
		},
		initialMessages: [
			{
				id: 'system',
				content: `For the full conversation, user have provided the following images: ${mediaList.map(media => media.name)}. Please help reply to user regarding these images`,
				dataset: mediaList,
				role: 'system',
			},
		] as MessageWithSelectedDataset[],
	});

	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;

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

export default useChatWithMedia;