diff --git a/web/containers/Providers/EventHandler.tsx b/web/containers/Providers/EventHandler.tsx index f871331b9..e1e0ff2a8 100644 --- a/web/containers/Providers/EventHandler.tsx +++ b/web/containers/Providers/EventHandler.tsx @@ -177,7 +177,6 @@ export default function EventHandler({ children }: { children: ReactNode }) { ) if (message.status === MessageStatus.Pending) { if (message.content.length) { - updateThreadWaiting(message.thread_id, false) setIsGeneratingResponse(false) } return diff --git a/web/helpers/atoms/ChatMessage.atom.ts b/web/helpers/atoms/ChatMessage.atom.ts index 904ad344d..0061eca7c 100644 --- a/web/helpers/atoms/ChatMessage.atom.ts +++ b/web/helpers/atoms/ChatMessage.atom.ts @@ -16,6 +16,8 @@ import { */ export const chatMessages = atom>({}) +export const readyThreadsMessagesAtom = atom>({}) + /** * Return the chat messages for the current active conversation */ @@ -34,6 +36,10 @@ export const setConvoMessagesAtom = atom( } newData[threadId] = messages set(chatMessages, newData) + set(readyThreadsMessagesAtom, { + ...get(readyThreadsMessagesAtom), + [threadId]: true, + }) } ) diff --git a/web/hooks/useCreateNewThread.ts b/web/hooks/useCreateNewThread.ts index 14f76306f..722e5b7e4 100644 --- a/web/hooks/useCreateNewThread.ts +++ b/web/hooks/useCreateNewThread.ts @@ -9,6 +9,8 @@ import { ThreadState, Model, AssistantTool, + events, + InferenceEvent, } from '@janhq/core' import { atom, useAtomValue, useSetAtom } from 'jotai' @@ -30,6 +32,7 @@ import { threadStatesAtom, updateThreadAtom, setThreadModelParamsAtom, + isGeneratingResponseAtom, } from '@/helpers/atoms/Thread.atom' const createNewThreadAtom = atom(null, (get, set, newThread: Thread) => { @@ -57,6 +60,7 @@ export const useCreateNewThread = () => { const setSelectedModel = useSetAtom(selectedModelAtom) const setThreadModelParams = useSetAtom(setThreadModelParamsAtom) const { experimentalFeature } = useContext(FeatureToggleContext) + const setIsGeneratingResponse = useSetAtom(isGeneratingResponseAtom) const { recommendedModel, downloadedModels } = useRecommendedModel() @@ -66,6 +70,10 @@ export const useCreateNewThread = () => { assistant: Assistant, model?: Model | undefined ) => { + // Stop generating if any + setIsGeneratingResponse(false) + events.emit(InferenceEvent.OnInferenceStopped, {}) + const defaultModel = model ?? recommendedModel ?? downloadedModels[0] // check last thread message, if there empty last message use can not create thread diff --git a/web/hooks/useSetActiveThread.ts b/web/hooks/useSetActiveThread.ts index 4bcd223eb..486a14d03 100644 --- a/web/hooks/useSetActiveThread.ts +++ b/web/hooks/useSetActiveThread.ts @@ -1,20 +1,14 @@ -import { useCallback } from 'react' +import { ExtensionTypeEnum, Thread, ConversationalExtension } from '@janhq/core' -import { - InferenceEvent, - ExtensionTypeEnum, - Thread, - events, - ConversationalExtension, -} from '@janhq/core' - -import { useSetAtom } from 'jotai' +import { useAtomValue, useSetAtom } from 'jotai' import { extensionManager } from '@/extension' -import { setConvoMessagesAtom } from '@/helpers/atoms/ChatMessage.atom' +import { + readyThreadsMessagesAtom, + setConvoMessagesAtom, +} from '@/helpers/atoms/ChatMessage.atom' import { ModelParams, - isGeneratingResponseAtom, setActiveThreadIdAtom, setThreadModelParamsAtom, } from '@/helpers/atoms/Thread.atom' @@ -23,31 +17,22 @@ export default function useSetActiveThread() { const setActiveThreadId = useSetAtom(setActiveThreadIdAtom) const setThreadMessage = useSetAtom(setConvoMessagesAtom) const setThreadModelParams = useSetAtom(setThreadModelParamsAtom) - const setIsGeneratingResponse = useSetAtom(isGeneratingResponseAtom) + const readyMessageThreads = useAtomValue(readyThreadsMessagesAtom) - const setActiveThread = useCallback( - async (thread: Thread) => { - setIsGeneratingResponse(false) - events.emit(InferenceEvent.OnInferenceStopped, thread.id) - - // load the corresponding messages + const setActiveThread = async (thread: Thread) => { + // Load local messages only if there are no messages in the state + if (!readyMessageThreads[thread.id]) { const messages = await getLocalThreadMessage(thread.id) setThreadMessage(thread.id, messages) + } - setActiveThreadId(thread.id) - const modelParams: ModelParams = { - ...thread.assistants[0]?.model?.parameters, - ...thread.assistants[0]?.model?.settings, - } - setThreadModelParams(thread.id, modelParams) - }, - [ - setActiveThreadId, - setThreadMessage, - setThreadModelParams, - setIsGeneratingResponse, - ] - ) + setActiveThreadId(thread.id) + const modelParams: ModelParams = { + ...thread.assistants[0]?.model?.parameters, + ...thread.assistants[0]?.model?.settings, + } + setThreadModelParams(thread.id, modelParams) + } return { setActiveThread } } diff --git a/web/hooks/useThreads.ts b/web/hooks/useThreads.ts index 74bbb8c91..fd0b3456d 100644 --- a/web/hooks/useThreads.ts +++ b/web/hooks/useThreads.ts @@ -9,8 +9,6 @@ import { import { useSetAtom } from 'jotai' -import useSetActiveThread from './useSetActiveThread' - import { extensionManager } from '@/extension/ExtensionManager' import { ModelParams, @@ -24,7 +22,6 @@ const useThreads = () => { const setThreadStates = useSetAtom(threadStatesAtom) const setThreads = useSetAtom(threadsAtom) const setThreadModelRuntimeParams = useSetAtom(threadModelParamsAtom) - const { setActiveThread } = useSetActiveThread() const setThreadDataReady = useSetAtom(threadDataReadyAtom) useEffect(() => { @@ -56,16 +53,11 @@ const useThreads = () => { setThreadStates(localThreadStates) setThreads(localThreads) setThreadModelRuntimeParams(threadModelParams) - - if (localThreads.length > 0) { - setActiveThread(localThreads[0]) - } setThreadDataReady(true) } getThreads() }, [ - setActiveThread, setThreadModelRuntimeParams, setThreadStates, setThreads, diff --git a/web/screens/Chat/ChatInput/index.tsx b/web/screens/Chat/ChatInput/index.tsx index b53c4d651..f6bf60ba4 100644 --- a/web/screens/Chat/ChatInput/index.tsx +++ b/web/screens/Chat/ChatInput/index.tsx @@ -38,6 +38,8 @@ import { getCurrentChatMessagesAtom } from '@/helpers/atoms/ChatMessage.atom' import { activeThreadAtom, getActiveThreadIdAtom, + isGeneratingResponseAtom, + threadStatesAtom, waitingToSendMessage, } from '@/helpers/atoms/Thread.atom' @@ -57,6 +59,12 @@ const ChatInput: React.FC = () => { const imageInputRef = useRef(null) const [showAttacmentMenus, setShowAttacmentMenus] = useState(false) const { experimentalFeature } = useContext(FeatureToggleContext) + const isGeneratingResponse = useAtomValue(isGeneratingResponseAtom) + const threadStates = useAtomValue(threadStatesAtom) + + const isStreamingResponse = Object.values(threadStates).some( + (threadState) => threadState.waitingForResponse + ) const onPromptChange = (e: React.ChangeEvent) => { setCurrentPrompt(e.target.value) @@ -235,7 +243,9 @@ const ChatInput: React.FC = () => { accept="application/pdf" /> - {messages[messages.length - 1]?.status !== MessageStatus.Pending ? ( + {messages[messages.length - 1]?.status !== MessageStatus.Pending && + !isGeneratingResponse && + !isStreamingResponse ? (