From 8fc78086547e414358dbbeb1d4f9d286fa2d0618 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Wed, 20 Aug 2025 20:20:50 +0530 Subject: [PATCH] feat(hooks): implement useChat hook --- src/lib/hooks/useChat.tsx | 643 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 643 insertions(+) create mode 100644 src/lib/hooks/useChat.tsx diff --git a/src/lib/hooks/useChat.tsx b/src/lib/hooks/useChat.tsx new file mode 100644 index 0000000..573ac6b --- /dev/null +++ b/src/lib/hooks/useChat.tsx @@ -0,0 +1,643 @@ +'use client'; + +import { Message } from '@/components/ChatWindow'; +import { createContext, useContext, useEffect, useRef, useState } from 'react'; +import crypto from 'crypto'; +import { useSearchParams } from 'next/navigation'; +import { toast } from 'sonner'; +import { Document } from '@langchain/core/documents'; +import { getSuggestions } from '../actions'; + +type ChatContext = { + messages: Message[]; + chatHistory: [string, string][]; + files: File[]; + fileIds: string[]; + focusMode: string; + chatId: string | undefined; + optimizationMode: string; + isMessagesLoaded: boolean; + loading: boolean; + notFound: boolean; + messageAppeared: boolean; + isReady: boolean; + hasError: boolean; + setOptimizationMode: (mode: string) => void; + setFocusMode: (mode: string) => void; + setFiles: (files: File[]) => void; + setFileIds: (fileIds: string[]) => void; + sendMessage: ( + message: string, + messageId?: string, + rewrite?: boolean, + ) => Promise; + rewrite: (messageId: string) => void; +}; + +export interface File { + fileName: string; + fileExtension: string; + fileId: string; +} + +interface ChatModelProvider { + name: string; + provider: string; +} + +interface EmbeddingModelProvider { + name: string; + provider: string; +} + +const checkConfig = async ( + setChatModelProvider: (provider: ChatModelProvider) => void, + setEmbeddingModelProvider: (provider: EmbeddingModelProvider) => void, + setIsConfigReady: (ready: boolean) => void, + setHasError: (hasError: boolean) => void, +) => { + try { + let chatModel = localStorage.getItem('chatModel'); + let chatModelProvider = localStorage.getItem('chatModelProvider'); + let embeddingModel = localStorage.getItem('embeddingModel'); + let embeddingModelProvider = localStorage.getItem('embeddingModelProvider'); + + const autoImageSearch = localStorage.getItem('autoImageSearch'); + const autoVideoSearch = localStorage.getItem('autoVideoSearch'); + + if (!autoImageSearch) { + localStorage.setItem('autoImageSearch', 'true'); + } + + if (!autoVideoSearch) { + localStorage.setItem('autoVideoSearch', 'false'); + } + + const providers = await fetch(`/api/models`, { + headers: { + 'Content-Type': 'application/json', + }, + }).then(async (res) => { + if (!res.ok) + throw new Error( + `Failed to fetch models: ${res.status} ${res.statusText}`, + ); + return res.json(); + }); + + if ( + !chatModel || + !chatModelProvider || + !embeddingModel || + !embeddingModelProvider + ) { + if (!chatModel || !chatModelProvider) { + const chatModelProviders = providers.chatModelProviders; + const chatModelProvidersKeys = Object.keys(chatModelProviders); + + if (!chatModelProviders || chatModelProvidersKeys.length === 0) { + return toast.error('No chat models available'); + } else { + chatModelProvider = + chatModelProvidersKeys.find( + (provider) => + Object.keys(chatModelProviders[provider]).length > 0, + ) || chatModelProvidersKeys[0]; + } + + if ( + chatModelProvider === 'custom_openai' && + Object.keys(chatModelProviders[chatModelProvider]).length === 0 + ) { + toast.error( + "Looks like you haven't configured any chat model providers. Please configure them from the settings page or the config file.", + ); + return setHasError(true); + } + + chatModel = Object.keys(chatModelProviders[chatModelProvider])[0]; + } + + if (!embeddingModel || !embeddingModelProvider) { + const embeddingModelProviders = providers.embeddingModelProviders; + + if ( + !embeddingModelProviders || + Object.keys(embeddingModelProviders).length === 0 + ) + return toast.error('No embedding models available'); + + embeddingModelProvider = Object.keys(embeddingModelProviders)[0]; + embeddingModel = Object.keys( + embeddingModelProviders[embeddingModelProvider], + )[0]; + } + + localStorage.setItem('chatModel', chatModel!); + localStorage.setItem('chatModelProvider', chatModelProvider); + localStorage.setItem('embeddingModel', embeddingModel!); + localStorage.setItem('embeddingModelProvider', embeddingModelProvider); + } else { + const chatModelProviders = providers.chatModelProviders; + const embeddingModelProviders = providers.embeddingModelProviders; + + if ( + Object.keys(chatModelProviders).length > 0 && + (!chatModelProviders[chatModelProvider] || + Object.keys(chatModelProviders[chatModelProvider]).length === 0) + ) { + const chatModelProvidersKeys = Object.keys(chatModelProviders); + chatModelProvider = + chatModelProvidersKeys.find( + (key) => Object.keys(chatModelProviders[key]).length > 0, + ) || chatModelProvidersKeys[0]; + + localStorage.setItem('chatModelProvider', chatModelProvider); + } + + if ( + chatModelProvider && + !chatModelProviders[chatModelProvider][chatModel] + ) { + if ( + chatModelProvider === 'custom_openai' && + Object.keys(chatModelProviders[chatModelProvider]).length === 0 + ) { + toast.error( + "Looks like you haven't configured any chat model providers. Please configure them from the settings page or the config file.", + ); + return setHasError(true); + } + + chatModel = Object.keys( + chatModelProviders[ + Object.keys(chatModelProviders[chatModelProvider]).length > 0 + ? chatModelProvider + : Object.keys(chatModelProviders)[0] + ], + )[0]; + + localStorage.setItem('chatModel', chatModel); + } + + if ( + Object.keys(embeddingModelProviders).length > 0 && + !embeddingModelProviders[embeddingModelProvider] + ) { + embeddingModelProvider = Object.keys(embeddingModelProviders)[0]; + localStorage.setItem('embeddingModelProvider', embeddingModelProvider); + } + + if ( + embeddingModelProvider && + !embeddingModelProviders[embeddingModelProvider][embeddingModel] + ) { + embeddingModel = Object.keys( + embeddingModelProviders[embeddingModelProvider], + )[0]; + localStorage.setItem('embeddingModel', embeddingModel); + } + } + + setChatModelProvider({ + name: chatModel!, + provider: chatModelProvider, + }); + + setEmbeddingModelProvider({ + name: embeddingModel!, + provider: embeddingModelProvider, + }); + + setIsConfigReady(true); + } catch (err) { + console.error('An error occurred while checking the configuration:', err); + setIsConfigReady(false); + setHasError(true); + } +}; + +const loadMessages = async ( + chatId: string, + setMessages: (messages: Message[]) => void, + setIsMessagesLoaded: (loaded: boolean) => void, + setChatHistory: (history: [string, string][]) => void, + setFocusMode: (mode: string) => void, + setNotFound: (notFound: boolean) => void, + setFiles: (files: File[]) => void, + setFileIds: (fileIds: string[]) => void, +) => { + const res = await fetch(`/api/chats/${chatId}`, { + method: 'GET', + headers: { + 'Content-Type': 'application/json', + }, + }); + + if (res.status === 404) { + setNotFound(true); + setIsMessagesLoaded(true); + return; + } + + const data = await res.json(); + + const messages = data.messages.map((msg: any) => { + return { + ...msg, + ...JSON.parse(msg.metadata), + }; + }) as Message[]; + + setMessages(messages); + + const history = messages.map((msg) => { + return [msg.role, msg.content]; + }) as [string, string][]; + + console.debug(new Date(), 'app:messages_loaded'); + + document.title = messages[0].content; + + const files = data.chat.files.map((file: any) => { + return { + fileName: file.name, + fileExtension: file.name.split('.').pop(), + fileId: file.fileId, + }; + }); + + setFiles(files); + setFileIds(files.map((file: File) => file.fileId)); + + setChatHistory(history); + setFocusMode(data.chat.focusMode); + setIsMessagesLoaded(true); +}; + +export const chatContext = createContext({ + chatHistory: [], + chatId: '', + fileIds: [], + files: [], + focusMode: '', + hasError: false, + isMessagesLoaded: false, + isReady: false, + loading: false, + messageAppeared: false, + messages: [], + notFound: false, + optimizationMode: '', + rewrite: () => {}, + sendMessage: async () => {}, + setFileIds: () => {}, + setFiles: () => {}, + setFocusMode: () => {}, + setOptimizationMode: () => {}, +}); + +export const ChatProvider = ({ + children, + id, +}: { + children: React.ReactNode; + id?: string; +}) => { + const searchParams = useSearchParams(); + const initialMessage = searchParams.get('q'); + + const [chatId, setChatId] = useState(id); + const [newChatCreated, setNewChatCreated] = useState(false); + + const [loading, setLoading] = useState(false); + const [messageAppeared, setMessageAppeared] = useState(false); + + const [chatHistory, setChatHistory] = useState<[string, string][]>([]); + const [messages, setMessages] = useState([]); + + const [files, setFiles] = useState([]); + const [fileIds, setFileIds] = useState([]); + + const [focusMode, setFocusMode] = useState('webSearch'); + const [optimizationMode, setOptimizationMode] = useState('speed'); + + const [isMessagesLoaded, setIsMessagesLoaded] = useState(false); + + const [notFound, setNotFound] = useState(false); + + const [chatModelProvider, setChatModelProvider] = useState( + { + name: '', + provider: '', + }, + ); + + const [embeddingModelProvider, setEmbeddingModelProvider] = + useState({ + name: '', + provider: '', + }); + + const [isConfigReady, setIsConfigReady] = useState(false); + const [hasError, setHasError] = useState(false); + const [isReady, setIsReady] = useState(false); + + const messagesRef = useRef([]); + + useEffect(() => { + checkConfig( + setChatModelProvider, + setEmbeddingModelProvider, + setIsConfigReady, + setHasError, + ); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []); + + useEffect(() => { + if ( + chatId && + !newChatCreated && + !isMessagesLoaded && + messages.length === 0 + ) { + loadMessages( + chatId, + setMessages, + setIsMessagesLoaded, + setChatHistory, + setFocusMode, + setNotFound, + setFiles, + setFileIds, + ); + } else if (!chatId) { + setNewChatCreated(true); + setIsMessagesLoaded(true); + setChatId(crypto.randomBytes(20).toString('hex')); + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []); + + useEffect(() => { + messagesRef.current = messages; + }, [messages]); + + useEffect(() => { + if (isMessagesLoaded && isConfigReady) { + setIsReady(true); + console.debug(new Date(), 'app:ready'); + } else { + setIsReady(false); + } + }, [isMessagesLoaded, isConfigReady]); + + const rewrite = (messageId: string) => { + const index = messages.findIndex((msg) => msg.messageId === messageId); + + if (index === -1) return; + + const message = messages[index - 1]; + + setMessages((prev) => { + return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)]; + }); + setChatHistory((prev) => { + return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)]; + }); + + sendMessage(message.content, message.messageId, true); + }; + + useEffect(() => { + if (isReady && initialMessage && isConfigReady) { + if (!isConfigReady) { + toast.error('Cannot send message before the configuration is ready'); + return; + } + sendMessage(initialMessage); + } + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [isConfigReady, isReady, initialMessage]); + + const sendMessage: ChatContext['sendMessage'] = async ( + message, + messageId, + rewrite = false, + ) => { + if (loading) return; + setLoading(true); + setMessageAppeared(false); + + let sources: Document[] | undefined = undefined; + let recievedMessage = ''; + let added = false; + + messageId = messageId ?? crypto.randomBytes(7).toString('hex'); + + setMessages((prevMessages) => [ + ...prevMessages, + { + content: message, + messageId: messageId, + chatId: chatId!, + role: 'user', + createdAt: new Date(), + }, + ]); + + const messageHandler = async (data: any) => { + if (data.type === 'error') { + toast.error(data.data); + setLoading(false); + return; + } + + if (data.type === 'sources') { + sources = data.data; + if (!added) { + setMessages((prevMessages) => [ + ...prevMessages, + { + content: '', + messageId: data.messageId, + chatId: chatId!, + role: 'assistant', + sources: sources, + createdAt: new Date(), + }, + ]); + added = true; + } + setMessageAppeared(true); + } + + if (data.type === 'message') { + if (!added) { + setMessages((prevMessages) => [ + ...prevMessages, + { + content: data.data, + messageId: data.messageId, + chatId: chatId!, + role: 'assistant', + sources: sources, + createdAt: new Date(), + }, + ]); + added = true; + } + + setMessages((prev) => + prev.map((message) => { + if (message.messageId === data.messageId) { + return { ...message, content: message.content + data.data }; + } + + return message; + }), + ); + + recievedMessage += data.data; + setMessageAppeared(true); + } + + if (data.type === 'messageEnd') { + setChatHistory((prevHistory) => [ + ...prevHistory, + ['human', message], + ['assistant', recievedMessage], + ]); + + setLoading(false); + + const lastMsg = messagesRef.current[messagesRef.current.length - 1]; + + const autoImageSearch = localStorage.getItem('autoImageSearch'); + const autoVideoSearch = localStorage.getItem('autoVideoSearch'); + + if (autoImageSearch === 'true') { + document + .getElementById(`search-images-${lastMsg.messageId}`) + ?.click(); + } + + if (autoVideoSearch === 'true') { + document + .getElementById(`search-videos-${lastMsg.messageId}`) + ?.click(); + } + + if ( + lastMsg.role === 'assistant' && + lastMsg.sources && + lastMsg.sources.length > 0 && + !lastMsg.suggestions + ) { + const suggestions = await getSuggestions(messagesRef.current); + setMessages((prev) => + prev.map((msg) => { + if (msg.messageId === lastMsg.messageId) { + return { ...msg, suggestions: suggestions }; + } + return msg; + }), + ); + } + } + }; + + const messageIndex = messages.findIndex((m) => m.messageId === messageId); + + const res = await fetch('/api/chat', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + content: message, + message: { + messageId: messageId, + chatId: chatId!, + content: message, + }, + chatId: chatId!, + files: fileIds, + focusMode: focusMode, + optimizationMode: optimizationMode, + history: rewrite + ? chatHistory.slice(0, messageIndex === -1 ? undefined : messageIndex) + : chatHistory, + chatModel: { + name: chatModelProvider.name, + provider: chatModelProvider.provider, + }, + embeddingModel: { + name: embeddingModelProvider.name, + provider: embeddingModelProvider.provider, + }, + systemInstructions: localStorage.getItem('systemInstructions'), + }), + }); + + if (!res.body) throw new Error('No response body'); + + const reader = res.body?.getReader(); + const decoder = new TextDecoder('utf-8'); + + let partialChunk = ''; + + while (true) { + const { value, done } = await reader.read(); + if (done) break; + + partialChunk += decoder.decode(value, { stream: true }); + + try { + const messages = partialChunk.split('\n'); + for (const msg of messages) { + if (!msg.trim()) continue; + const json = JSON.parse(msg); + messageHandler(json); + } + partialChunk = ''; + } catch (error) { + console.warn('Incomplete JSON, waiting for next chunk...'); + } + } + }; + + return ( + + {children} + + ); +}; + +export const useChat = () => { + const ctx = useContext(chatContext); + return ctx; +};