fix: duplicated messages when user switch between conversations (#441)

Signed-off-by: James <james@jan.ai>
Co-authored-by: James <james@jan.ai>

electron/package.json

@@ -67,7 +67,6 @@
     "electron-store": "^8.1.0",
     "electron-updater": "^6.1.4",
     "pacote": "^17.0.4",
-    "react-intersection-observer": "^9.5.2",
     "request": "^2.88.2",
     "request-progress": "^3.0.0",
     "use-debounce": "^9.0.4"

web/app/_components/BasicPromptInput/index.tsx

@@ -4,7 +4,6 @@ import { currentPromptAtom } from '@helpers/JotaiWrapper'
 import { getActiveConvoIdAtom } from '@helpers/atoms/Conversation.atom'
 import { selectedModelAtom } from '@helpers/atoms/Model.atom'
 import useCreateConversation from '@hooks/useCreateConversation'
-import useInitModel from '@hooks/useInitModel'
 import useSendChatMessage from '@hooks/useSendChatMessage'
 import { useAtom, useAtomValue } from 'jotai'
 import { ChangeEvent, useEffect, useRef } from 'react'
@@ -16,8 +15,6 @@ const BasicPromptInput: React.FC = () => {
   const { sendChatMessage } = useSendChatMessage()
   const { requestCreateConvo } = useCreateConversation()
 
-  const { initModel } = useInitModel()
-
   const textareaRef = useRef<HTMLTextAreaElement>(null)
 
   const handleKeyDown = async (
@@ -35,7 +32,6 @@ const BasicPromptInput: React.FC = () => {
           }
 
           await requestCreateConvo(selectedModel)
-          await initModel(selectedModel)
           sendChatMessage()
         }
       }

web/app/_components/ChatBody/index.tsx

@@ -1,60 +1,17 @@
 'use client'
 
-import React, { useCallback, useRef, useState, useEffect } from 'react'
+import React from 'react'
 import ChatItem from '../ChatItem'
 import useChatMessages from '@hooks/useChatMessages'
-import { useAtomValue } from 'jotai'
-import { selectAtom } from 'jotai/utils'
-import { getActiveConvoIdAtom } from '@helpers/atoms/Conversation.atom'
-import { chatMessages } from '@helpers/atoms/ChatMessage.atom'
 
 const ChatBody: React.FC = () => {
-  const activeConversationId = useAtomValue(getActiveConvoIdAtom) ?? ''
-  const messageList = useAtomValue(
-    selectAtom(
-      chatMessages,
-      useCallback((v) => v[activeConversationId], [activeConversationId])
-    )
-  )
-  const [content, setContent] = useState<React.JSX.Element[]>([])
-
-  const [offset, setOffset] = useState(0)
-  const { loading, hasMore } = useChatMessages(offset)
-  const intersectObs = useRef<any>(null)
-
-  const lastPostRef = useCallback(
-    (message: ChatMessage) => {
-      if (loading) return
-
-      if (intersectObs.current) intersectObs.current.disconnect()
-
-      intersectObs.current = new IntersectionObserver((entries) => {
-        if (entries[0].isIntersecting && hasMore) {
-          setOffset((prevOffset) => prevOffset + 5)
-        }
-      })
-
-      if (message) intersectObs.current.observe(message)
-    },
-    [loading, hasMore]
-  )
-
-  useEffect(() => {
-    const list = messageList?.map((message, index) => {
-      if (messageList?.length === index + 1) {
-        return (
-          // @ts-ignore
-          <ChatItem ref={lastPostRef} message={message} key={message.id} />
-        )
-      }
-      return <ChatItem message={message} key={message.id} />
-    })
-    setContent(list)
-  }, [messageList, lastPostRef])
+  const { messages } = useChatMessages()
 
   return (
-    <div className="[&>*:nth-child(odd)]:bg-background flex h-full flex-1 flex-col-reverse overflow-y-auto">
-      {content}
+    <div className="flex h-full flex-1 flex-col-reverse overflow-y-auto [&>*:nth-child(odd)]:bg-background">
+      {messages.map((message) => (
+        <ChatItem message={message} key={message.id} />
+      ))}
     </div>
   )
 }

web/app/_components/HistoryItem/index.tsx

@@ -1,19 +1,16 @@
 import React from 'react'
 import { useAtomValue, useSetAtom } from 'jotai'
-import Image from 'next/image'
 import { ModelManagementService } from '@janhq/core'
 import { executeSerial } from '../../../../electron/core/plugin-manager/execution/extension-manager'
 import {
   getActiveConvoIdAtom,
   setActiveConvoIdAtom,
-  updateConversationErrorAtom,
   updateConversationWaitingForResponseAtom,
 } from '@helpers/atoms/Conversation.atom'
 import {
   setMainViewStateAtom,
   MainViewState,
 } from '@helpers/atoms/MainView.atom'
-import useInitModel from '@hooks/useInitModel'
 import { displayDate } from '@utils/datetime'
 import { twMerge } from 'tailwind-merge'
 
@@ -36,11 +33,8 @@ const HistoryItem: React.FC<Props> = ({
   const activeConvoId = useAtomValue(getActiveConvoIdAtom)
   const setActiveConvoId = useSetAtom(setActiveConvoIdAtom)
   const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
-  const updateConvError = useSetAtom(updateConversationErrorAtom)
   const isSelected = activeConvoId === conversation._id
 
-  const { initModel } = useInitModel()
-
   const onClick = async () => {
     const model = await executeSerial(
       ModelManagementService.GetModelById,
@@ -48,13 +42,6 @@ const HistoryItem: React.FC<Props> = ({
     )
 
     if (conversation._id) updateConvWaiting(conversation._id, true)
-    initModel(model).then((res: any) => {
-      if (conversation._id) updateConvWaiting(conversation._id, false)
-
-      if (res?.error && conversation._id) {
-        updateConvError(conversation._id, res.error)
-      }
-    })
 
     if (activeConvoId !== conversation._id) {
       setMainViewState(MainViewState.Conversation)

web/app/_components/InputToolbar/index.tsx

@@ -4,9 +4,8 @@
 import BasicPromptInput from '../BasicPromptInput'
 import BasicPromptAccessories from '../BasicPromptAccessories'
 import { useAtomValue, useSetAtom } from 'jotai'
-import { showingAdvancedPromptAtom } from '@helpers/atoms/Modal.atom'
 import SecondaryButton from '../SecondaryButton'
-import { Fragment, useEffect, useState } from 'react'
+import { useEffect, useState } from 'react'
 import { PlusIcon } from '@heroicons/react/24/outline'
 import useCreateConversation from '@hooks/useCreateConversation'
 import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
@@ -19,7 +18,6 @@ import { activeBotAtom } from '@helpers/atoms/Bot.atom'
 import { useGetDownloadedModels } from '@hooks/useGetDownloadedModels'
 
 const InputToolbar: React.FC = () => {
-  const showingAdvancedPrompt = useAtomValue(showingAdvancedPromptAtom)
   const activeModel = useAtomValue(activeAssistantModelAtom)
   const { requestCreateConvo } = useCreateConversation()
   const currentConvoState = useAtomValue(currentConvoStateAtom)
@@ -76,7 +74,7 @@ const InputToolbar: React.FC = () => {
 
   if (inputState === 'disabled')
     return (
-      <div className="bg-background/90 sticky bottom-0 flex items-center justify-center">
+      <div className="sticky bottom-0 flex items-center justify-center bg-background/90">
         <p className="mx-auto my-5 line-clamp-2 text-ellipsis text-center italic text-gray-600">
           {error}
         </p>
@@ -84,7 +82,7 @@ const InputToolbar: React.FC = () => {
     )
 
   return (
-    <div className="bg-background/90 sticky bottom-0 w-full px-5 py-0">
+    <div className="sticky bottom-0 w-full bg-background/90 px-5 py-0">
       {currentConvoState?.error && (
         <div className="flex flex-row justify-center">
           <span className="mx-5 my-2 text-sm text-red-500">

web/app/_components/ModelActionButton/index.tsx

@@ -24,12 +24,14 @@ const modelActionMapper: Record<ModelActionType, ModelActionStyle> = {
 }
 
 type Props = {
+  disabled?: boolean
   type: ModelActionType
   onActionClick: (type: ModelActionType) => void
   onDeleteClick: () => void
 }
 
 const ModelActionButton: React.FC<Props> = ({
+  disabled = false,
   type,
   onActionClick,
   onDeleteClick,
@@ -48,6 +50,7 @@ const ModelActionButton: React.FC<Props> = ({
       <div className="flex items-center justify-end gap-x-4">
         <ModelActionMenu onDeleteClick={onDeleteClick} />
         <Button
+          disabled={disabled}
           size="sm"
           themes={styles.title === 'Start' ? 'accent' : 'default'}
           onClick={() => onClick()}

web/app/_components/ModelRow/index.tsx

@@ -12,7 +12,7 @@ type Props = {
 }
 
 const ModelRow: React.FC<Props> = ({ model }) => {
-  const { startModel, stopModel } = useStartStopModel()
+  const { loading, startModel, stopModel } = useStartStopModel()
   const activeModel = useAtomValue(activeAssistantModelAtom)
   const { deleteModel } = useDeleteModel()
 
@@ -57,6 +57,7 @@ const ModelRow: React.FC<Props> = ({ model }) => {
         <ModelStatusComponent status={status} />
       </td>
       <ModelActionButton
+        disabled={loading}
         type={actionButtonType}
         onActionClick={onModelActionClick}
         onDeleteClick={onDeleteClick}

web/app/_components/SidebarEmptyHistory/index.tsx

@@ -1,5 +1,4 @@
 import useCreateConversation from '@hooks/useCreateConversation'
-import PrimaryButton from '../PrimaryButton'
 import { useAtomValue, useSetAtom } from 'jotai'
 import { useEffect, useState } from 'react'
 import {
@@ -7,11 +6,9 @@ import {
   setMainViewStateAtom,
 } from '@helpers/atoms/MainView.atom'
 import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
-import useInitModel from '@hooks/useInitModel'
 import { useGetDownloadedModels } from '@hooks/useGetDownloadedModels'
 import { Button } from '@uikit'
-
-import {MessageCircle} from "lucide-react"
+import { MessageCircle } from 'lucide-react'
 
 enum ActionButton {
   DownloadModel = 'Download a Model',
@@ -25,8 +22,6 @@ const SidebarEmptyHistory: React.FC = () => {
   const { requestCreateConvo } = useCreateConversation()
   const [action, setAction] = useState(ActionButton.DownloadModel)
 
-  const { initModel } = useInitModel()
-
   useEffect(() => {
     if (downloadedModels.length > 0) {
       setAction(ActionButton.StartChat)
@@ -35,32 +30,27 @@ const SidebarEmptyHistory: React.FC = () => {
     }
   }, [downloadedModels])
 
-  const onClick = () => {
+  const onClick = async () => {
     if (action === ActionButton.DownloadModel) {
       setMainView(MainViewState.ExploreModel)
     } else {
       if (!activeModel) {
         setMainView(MainViewState.ConversationEmptyModel)
       } else {
-        createConversationAndInitModel(activeModel)
+        await requestCreateConvo(activeModel)
       }
     }
   }
 
-  const createConversationAndInitModel = async (model: AssistantModel) => {
-    await requestCreateConvo(model)
-    await initModel(model)
-  }
-
   return (
     <div className="flex flex-col items-center gap-3 py-10">
       <MessageCircle size={32} />
       <div className="flex flex-col items-center gap-y-2">
         <h6 className="text-center text-base">No Chat History</h6>
-        <p className="text-center text-muted-foreground mb-6">
+        <p className="mb-6 text-center text-muted-foreground">
           Get started by creating a new chat.
         </p>
-        <Button onClick={onClick} themes="accent" >
+        <Button onClick={onClick} themes="accent">
           {action}
         </Button>
       </div>

web/helpers/atoms/ChatMessage.atom.ts

@@ -4,14 +4,34 @@ import { getActiveConvoIdAtom } from './Conversation.atom'
 /**
  * Stores all chat messages for all conversations
  */
-export const chatMessages = atom<Record<string, ChatMessage[]>>({})
+const chatMessages = atom<Record<string, ChatMessage[]>>({})
 
-export const currentChatMessagesAtom = atom<ChatMessage[]>((get) => {
+/**
+ * Return the chat messages for the current active conversation
+ */
+export const getCurrentChatMessagesAtom = atom<ChatMessage[]>((get) => {
   const activeConversationId = get(getActiveConvoIdAtom)
   if (!activeConversationId) return []
   return get(chatMessages)[activeConversationId] ?? []
 })
 
+export const setCurrentChatMessagesAtom = atom(
+  null,
+  (get, set, messages: ChatMessage[]) => {
+    const currentConvoId = get(getActiveConvoIdAtom)
+    if (!currentConvoId) return
+
+    const newData: Record<string, ChatMessage[]> = {
+      ...get(chatMessages),
+    }
+    newData[currentConvoId] = messages
+    set(chatMessages, newData)
+  }
+)
+
+/**
+ * Used for pagination. Add old messages to the current conversation
+ */
 export const addOldMessagesAtom = atom(
   null,
   (get, set, newMessages: ChatMessage[]) => {

web/hooks/useChatMessages.ts

@@ -1,62 +1,50 @@
 import { toChatMessage } from '@models/ChatMessage'
 import { executeSerial } from '@services/pluginService'
 import { useAtomValue, useSetAtom } from 'jotai'
-import { useEffect, useState } from 'react'
+import { useEffect } from 'react'
 import { DataService } from '@janhq/core'
-import { addOldMessagesAtom } from '@helpers/atoms/ChatMessage.atom'
+import { getActiveConvoIdAtom } from '@helpers/atoms/Conversation.atom'
 import {
-  currentConversationAtom,
-  conversationStatesAtom,
-  updateConversationHasMoreAtom,
-} from '@helpers/atoms/Conversation.atom'
+  getCurrentChatMessagesAtom,
+  setCurrentChatMessagesAtom,
+} from '@helpers/atoms/ChatMessage.atom'
 
 /**
  * Custom hooks to get chat messages for current(active) conversation
- *
- * @param offset for pagination purpose
- * @returns
  */
-const useChatMessages = (offset = 0) => {
-  const [loading, setLoading] = useState(true)
-  const addOldChatMessages = useSetAtom(addOldMessagesAtom)
-  const currentConvo = useAtomValue(currentConversationAtom)
-  const convoStates = useAtomValue(conversationStatesAtom)
-  const updateConvoHasMore = useSetAtom(updateConversationHasMoreAtom)
+const useChatMessages = () => {
+  const setMessages = useSetAtom(setCurrentChatMessagesAtom)
+  const messages = useAtomValue(getCurrentChatMessagesAtom)
+  const activeConvoId = useAtomValue(getActiveConvoIdAtom)
+
+  const getMessages = async (convoId: string) => {
+    const data: any = await executeSerial(
+      DataService.GetConversationMessages,
+      convoId
+    )
+    if (!data) {
+      return []
+    }
+
+    return parseMessages(data)
+  }
 
   useEffect(() => {
-    if (!currentConvo) {
+    if (!activeConvoId) {
+      console.error('active convo is undefined')
       return
     }
-    const hasMore = convoStates[currentConvo._id ?? '']?.hasMore ?? true
-    if (!hasMore) return
 
-    const getMessages = async () => {
-      executeSerial(DataService.GetConversationMessages, currentConvo._id).then(
-        (data: any) => {
-          if (!data) {
-            return
-          }
-          const newMessages = parseMessages(data ?? [])
-          addOldChatMessages(newMessages)
-          updateConvoHasMore(currentConvo._id ?? '', false)
-          setLoading(false)
-        }
-      )
-    }
-    getMessages()
-  }, [
-    offset,
-    convoStates,
-    addOldChatMessages,
-    updateConvoHasMore,
-    currentConvo,
-  ])
+    getMessages(activeConvoId)
+      .then((messages) => {
+        setMessages(messages)
+      })
+      .catch((err) => {
+        console.error(err)
+      })
+  }, [activeConvoId])
 
-  return {
-    loading: loading,
-    error: undefined,
-    hasMore: convoStates[currentConvo?._id ?? '']?.hasMore ?? true,
-  }
+  return { messages }
 }
 
 function parseMessages(messages: RawMessage[]): ChatMessage[] {

web/hooks/useCreateConversation.ts

@@ -1,25 +1,18 @@
 import { useAtom, useSetAtom } from 'jotai'
-
 import { executeSerial } from '@services/pluginService'
 import { DataService, ModelManagementService } from '@janhq/core'
 import {
   userConversationsAtom,
   setActiveConvoIdAtom,
   addNewConversationStateAtom,
-  updateConversationWaitingForResponseAtom,
-  updateConversationErrorAtom,
 } from '@helpers/atoms/Conversation.atom'
-import useInitModel from './useInitModel'
 
 const useCreateConversation = () => {
-  const { initModel } = useInitModel()
   const [userConversations, setUserConversations] = useAtom(
     userConversationsAtom
   )
   const setActiveConvoId = useSetAtom(setActiveConvoIdAtom)
   const addNewConvoState = useSetAtom(addNewConversationStateAtom)
-  const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
-  const updateConvError = useSetAtom(updateConversationErrorAtom)
 
   const createConvoByBot = async (bot: Bot) => {
     const model = await executeSerial(
@@ -48,14 +41,6 @@ const useCreateConversation = () => {
     }
     const id = await executeSerial(DataService.CreateConversation, conv)
 
-    if (id) updateConvWaiting(id, true)
-    initModel(model).then((res: any) => {
-      if (id) updateConvWaiting(id, false)
-      if (res?.error) {
-        updateConvError(id, res.error)
-      }
-    })
-
     const mappedConvo: Conversation = {
       _id: id,
       modelId: model._id,

web/hooks/useDeleteConversation.ts

@@ -1,5 +1,5 @@
 import { currentPromptAtom } from '@helpers/JotaiWrapper'
-import { execute } from '@services/pluginService'
+import { executeSerial } from '@services/pluginService'
 import { useAtom, useAtomValue, useSetAtom } from 'jotai'
 import { DataService } from '@janhq/core'
 import { deleteConversationMessage } from '@helpers/atoms/ChatMessage.atom'
@@ -33,7 +33,7 @@ export default function useDeleteConversation() {
   const deleteConvo = async () => {
     if (activeConvoId) {
       try {
-        await execute(DataService.DeleteConversation, activeConvoId)
+        await executeSerial(DataService.DeleteConversation, activeConvoId)
         const currentConversations = userConversations.filter(
           (c) => c._id !== activeConvoId
         )

web/hooks/useInitModel.ts

@@ -1,32 +0,0 @@
-import { executeSerial } from '@services/pluginService'
-import { InferenceService } from '@janhq/core'
-import { useAtom } from 'jotai'
-import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
-
-export default function useInitModel() {
-  const [activeModel, setActiveModel] = useAtom(activeAssistantModelAtom)
-
-  const initModel = async (model: AssistantModel) => {
-    if (activeModel && activeModel._id === model._id) {
-      console.debug(`Model ${model._id} is already init. Ignore..`)
-      return
-    }
-
-    const currentTime = Date.now()
-    console.debug('Init model: ', model._id)
-
-    const res = await executeSerial(InferenceService.InitModel, model._id)
-    if (res?.error) {
-      console.error('Failed to init model: ', res.error)
-      return res
-    } else {
-      console.debug(
-        `Init model successfully!, take ${Date.now() - currentTime}ms`
-      )
-      setActiveModel(model)
-      return {}
-    }
-  }
-
-  return { initModel }
-}

web/hooks/useStartStopModel.ts

@@ -1,16 +1,22 @@
 import { executeSerial } from '@services/pluginService'
 import { ModelManagementService, InferenceService } from '@janhq/core'
-import useInitModel from './useInitModel'
-import { useSetAtom } from 'jotai'
+import { useAtom, useSetAtom } from 'jotai'
 import { activeAssistantModelAtom, stateModel } from '@helpers/atoms/Model.atom'
+import { useState } from 'react'
 
 export default function useStartStopModel() {
-  const { initModel } = useInitModel()
-  const setActiveModel = useSetAtom(activeAssistantModelAtom)
+  const [activeModel, setActiveModel] = useAtom(activeAssistantModelAtom)
+  const [loading, setLoading] = useState<boolean>(false)
   const setStateModel = useSetAtom(stateModel)
 
   const startModel = async (modelId: string) => {
+    if (activeModel && activeModel._id === modelId) {
+      console.debug(`Model ${modelId} is already init. Ignore..`)
+      return
+    }
+
     setStateModel({ state: 'start', loading: true, model: modelId })
+
     const model = await executeSerial(
       ModelManagementService.GetModelById,
       modelId
@@ -18,10 +24,23 @@ export default function useStartStopModel() {
     if (!model) {
       alert(`Model ${modelId} not found! Please re-download the model first.`)
       setStateModel((prev) => ({ ...prev, loading: false }))
-    } else {
-      await initModel(model)
-      setStateModel((prev) => ({ ...prev, loading: false }))
     }    
+    const currentTime = Date.now()
+    console.debug('Init model: ', model._id)
+
+    const res = await executeSerial(InferenceService.InitModel, model._id)
+
+    if (res?.error) {
+      const errorMessage = `Failed to init model: ${res.error}`
+      console.error(errorMessage)
+      alert(errorMessage)
+    } else {
+      console.debug(
+        `Init model successfully!, take ${Date.now() - currentTime}ms`
+      )
+      setActiveModel(model)
+    }
+    setLoading(false)
   }
 
   const stopModel = async (modelId: string) => {
@@ -33,5 +52,5 @@ export default function useStartStopModel() {
     }, 500)
   }
 
-  return { startModel, stopModel }
+  return { loading, startModel, stopModel }
 }

web/package.json

@@ -24,12 +24,11 @@
     "autoprefixer": "10.4.14",
     "class-variance-authority": "^0.7.0",
     "classnames": "^2.3.2",
-    "embla-carousel": "^8.0.0-rc11",
-    "embla-carousel-react": "^8.0.0-rc11",
     "eslint": "8.45.0",
     "eslint-config-next": "13.4.10",
     "framer-motion": "^10.16.4",
     "highlight.js": "^11.9.0",
+    "react-intersection-observer": "^9.5.2",
     "jotai": "^2.4.0",
     "jotai-optics": "^0.3.1",
     "jwt-decode": "^3.1.2",
@@ -39,7 +38,6 @@
     "next": "13.4.10",
     "next-auth": "^4.23.1",
     "next-themes": "^0.2.1",
-    "optics-ts": "^2.4.1",
     "postcss": "8.4.26",
     "react": "18.2.0",
     "react-dom": "18.2.0",