fix: download now change state immediately (janhq#475)

* fix: download now change state immediately Signed-off-by: James <[email protected]> * fix: add back last message Signed-off-by: James <[email protected]> * update send button status Signed-off-by: James <[email protected]> --------- Signed-off-by: James <[email protected]> Co-authored-by: James <[email protected]>
jesst3r · Oct 27, 2023 · 0bd52e5 · 0bd52e5
1 parent 7b7d2db
commit 0bd52e5
Show file tree

Hide file tree

Showing 9 changed files with 120 additions and 30 deletions.
diff --git a/core/events.ts b/core/events.ts
@@ -6,6 +6,7 @@ export enum EventName {
   OnNewMessageRequest = "onNewMessageRequest",
   OnNewMessageResponse = "onNewMessageResponse",
   OnMessageResponseUpdate = "onMessageResponseUpdate",
+  OnMessageResponseFinished = "OnMessageResponseFinished",
   OnDownloadUpdate = "onDownloadUpdate",
   OnDownloadSuccess = "onDownloadSuccess",
   OnDownloadError = "onDownloadError",

diff --git a/plugins/inference-plugin/index.ts b/plugins/inference-plugin/index.ts
@@ -140,6 +140,8 @@ async function handleMessageRequest(data: NewMessageRequest) {
       message.message = message.message.trim();
       // TODO: Common collections should be able to access via core functions instead of store
       await store.updateOne("messages", message._id, message);
+      events.emit("OnMessageResponseFinished", message);
+      // events.emit(EventName.OnMessageResponseFinished, message);
     },
     error: async (err) => {
       message.message =

diff --git a/web/app/_components/ExploreModelItemHeader/index.tsx b/web/app/_components/ExploreModelItemHeader/index.tsx
@@ -1,7 +1,5 @@
 import SimpleTag from '../SimpleTag'
-import PrimaryButton from '../PrimaryButton'
 import { formatDownloadPercentage, toGigabytes } from '@utils/converter'
-import SecondaryButton from '../SecondaryButton'
 import { useCallback, useEffect, useMemo } from 'react'
 import useGetPerformanceTag from '@hooks/useGetPerformanceTag'
 import useDownloadModel from '@hooks/useDownloadModel'
@@ -58,7 +56,6 @@ const ExploreModelItemHeader: React.FC<Props> = ({
   if (isDownloaded) {
     downloadButton = (
       <Button
-        size="sm"
         themes="accent"
         onClick={() => {
           setMainViewState(MainViewState.MyModel)
@@ -72,17 +69,20 @@ const ExploreModelItemHeader: React.FC<Props> = ({
   if (downloadState != null) {
     // downloading
     downloadButton = (
-      <SecondaryButton
+      <Button
         disabled
-        title={`Downloading (${formatDownloadPercentage(
-          downloadState.percent
-        )})`}
-      />
+        themes="accent"
+        onClick={() => {
+          setMainViewState(MainViewState.MyModel)
+        }}
+      >
+        Downloading {formatDownloadPercentage(downloadState.percent)}
+      </Button>
     )
   }
 
   return (
-    <div className="border-border bg-background/50 flex items-center justify-between rounded-t-md border-b px-4 py-2">
+    <div className="flex items-center justify-between rounded-t-md border-b border-border bg-background/50 px-4 py-2">
       <div className="flex items-center gap-2">
         <span>{exploreModel.name}</span>
         {performanceTag && (

diff --git a/web/app/_components/HistoryItem/index.tsx b/web/app/_components/HistoryItem/index.tsx
@@ -3,7 +3,6 @@ import { useAtomValue, useSetAtom } from 'jotai'
 import {
   getActiveConvoIdAtom,
   setActiveConvoIdAtom,
-  updateConversationWaitingForResponseAtom,
 } from '@helpers/atoms/Conversation.atom'
 import {
   setMainViewStateAtom,
@@ -12,7 +11,6 @@ import {
 import { displayDate } from '@utils/datetime'
 import { twMerge } from 'tailwind-merge'
 import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
-import { switchingModelConfirmationModalPropsAtom } from '@helpers/atoms/Modal.atom'
 import useStartStopModel from '@hooks/useStartStopModel'
 import useGetModelById from '@hooks/useGetModelById'
 
@@ -37,10 +35,6 @@ const HistoryItem: React.FC<Props> = ({
 
   const setMainViewState = useSetAtom(setMainViewStateAtom)
   const setActiveConvoId = useSetAtom(setActiveConvoIdAtom)
-  const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
-  const setConfirmationModalProps = useSetAtom(
-    switchingModelConfirmationModalPropsAtom
-  )
 
   const onClick = async () => {
     if (conversation.modelId == null) {
@@ -55,14 +49,13 @@ const HistoryItem: React.FC<Props> = ({
         startModel(model._id)
       } else if (activeModel._id !== model._id) {
         // display confirmation modal
-        setConfirmationModalProps({
-          replacingModel: model,
-        })
+        // TODO: temporarily disabled
+        // setConfirmationModalProps({
+        //   replacingModel: model,
+        // })
       }
     }
 
-    if (conversation._id) updateConvWaiting(conversation._id, true)
-
     if (activeConvoId !== conversation._id) {
       setMainViewState(MainViewState.Conversation)
       setActiveConvoId(conversation._id)

diff --git a/web/app/_components/SwitchingModelConfirmationModal/index.tsx b/web/app/_components/SwitchingModelConfirmationModal/index.tsx
@@ -82,11 +82,11 @@ const SwitchingModelConfirmationModal: React.FC = () => {
                       <p className="text-sm text-gray-500">
                         Selected conversation is using model{' '}
                         <span className="font-semibold text-black">
-                          {props?.replacingModel._id}
+                          {props?.replacingModel.name}
                         </span>
                         , but the active model is using{' '}
                         <span className="font-semibold text-black">
-                          {activeModel?._id}
+                          {activeModel?.name}
                         </span>
                         .
                       </p>
@@ -95,7 +95,7 @@ const SwitchingModelConfirmationModal: React.FC = () => {
                         Switch to
                         <span className="font-semibold text-black">
                           {' '}
-                          {props?.replacingModel._id}?
+                          {props?.replacingModel.name}?
                         </span>
                       </p>
                     </div>

diff --git a/web/helpers/EventHandler.tsx b/web/helpers/EventHandler.tsx
@@ -1,17 +1,35 @@
 import { addNewMessageAtom, updateMessageAtom } from './atoms/ChatMessage.atom'
 import { toChatMessage } from '@models/ChatMessage'
-import { events, EventName, NewMessageResponse } from '@janhq/core'
+import { events, EventName, NewMessageResponse, DataService } from '@janhq/core'
 import { useSetAtom } from 'jotai'
 import { ReactNode, useEffect } from 'react'
 import useGetBots from '@hooks/useGetBots'
 import useGetUserConversations from '@hooks/useGetUserConversations'
+import {
+  updateConversationAtom,
+  updateConversationWaitingForResponseAtom,
+} from './atoms/Conversation.atom'
+import { executeSerial } from '../../electron/core/plugin-manager/execution/extension-manager'
+import { debounce } from 'lodash'
+
+let currentConversation: Conversation | undefined = undefined
+
+const debouncedUpdateConversation = debounce(
+  async (updatedConv: Conversation) => {
+    await executeSerial(DataService.UpdateConversation, updatedConv)
+  },
+  1000
+)
 
 export default function EventHandler({ children }: { children: ReactNode }) {
   const addNewMessage = useSetAtom(addNewMessageAtom)
   const updateMessage = useSetAtom(updateMessageAtom)
+  const updateConversation = useSetAtom(updateConversationAtom)
   const { getBotById } = useGetBots()
   const { getConversationById } = useGetUserConversations()
 
+  const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
+
   async function handleNewMessageResponse(message: NewMessageResponse) {
     if (message.conversationId) {
       const convo = await getConversationById(message.conversationId)
@@ -34,25 +52,63 @@ export default function EventHandler({ children }: { children: ReactNode }) {
       messageResponse.conversationId &&
       messageResponse._id &&
       messageResponse.message
-    )
+    ) {
       updateMessage(
         messageResponse._id,
         messageResponse.conversationId,
         messageResponse.message
       )
+    }
+
+    if (messageResponse.conversationId) {
+      if (
+        !currentConversation ||
+        currentConversation._id !== messageResponse.conversationId
+      ) {
+        currentConversation = await getConversationById(
+          messageResponse.conversationId
+        )
+      }
+
+      const updatedConv: Conversation = {
+        ...currentConversation,
+        lastMessage: messageResponse.message,
+      }
+
+      updateConversation(updatedConv)
+      debouncedUpdateConversation(updatedConv)
+    }
+  }
+
+  async function handleMessageResponseFinished(
+    messageResponse: NewMessageResponse
+  ) {
+    if (!messageResponse.conversationId) return
+    console.debug('handleMessageResponseFinished', messageResponse)
+    updateConvWaiting(messageResponse.conversationId, false)
   }
 
   useEffect(() => {
     if (window.corePlugin.events) {
       events.on(EventName.OnNewMessageResponse, handleNewMessageResponse)
       events.on(EventName.OnMessageResponseUpdate, handleMessageResponseUpdate)
+      events.on(
+        "OnMessageResponseFinished",
+        // EventName.OnMessageResponseFinished,
+        handleMessageResponseFinished
+      )
     }
   }, [])
 
   useEffect(() => {
     return () => {
       events.off(EventName.OnNewMessageResponse, handleNewMessageResponse)
       events.off(EventName.OnMessageResponseUpdate, handleMessageResponseUpdate)
+      events.off(
+        "OnMessageResponseFinished",
+        // EventName.OnMessageResponseFinished,
+        handleMessageResponseFinished
+      )
     }
   }, [])
   return <>{children}</>

diff --git a/web/hooks/useDownloadModel.ts b/web/hooks/useDownloadModel.ts
@@ -1,7 +1,11 @@
 import { executeSerial } from '@services/pluginService'
-import { DataService, ModelManagementService } from '@janhq/core'
+import { ModelManagementService } from '@janhq/core'
+import { useSetAtom } from 'jotai'
+import { setDownloadStateAtom } from '@helpers/atoms/DownloadState.atom'
 
 export default function useDownloadModel() {
+  const setDownloadState = useSetAtom(setDownloadStateAtom)
+
   const assistanModel = (
     model: Product,
     modelVersion: ModelVersion
@@ -37,6 +41,22 @@ export default function useDownloadModel() {
   }
 
   const downloadModel = async (model: Product, modelVersion: ModelVersion) => {
+    // set an initial download state
+    setDownloadState({
+      modelId: modelVersion._id,
+      time: {
+        elapsed: 0,
+        remaining: 0,
+      },
+      speed: 0,
+      percent: 0,
+      size: {
+        total: 0,
+        transferred: 0,
+      },
+      fileName: modelVersion._id,
+    })
+
     modelVersion.startDownloadAt = Date.now()
     const assistantModel = assistanModel(model, modelVersion)
     await executeSerial(ModelManagementService.StoreModel, assistantModel)

diff --git a/web/hooks/useSendChatMessage.ts b/web/hooks/useSendChatMessage.ts
@@ -13,13 +13,14 @@ import { addNewMessageAtom } from '@helpers/atoms/ChatMessage.atom'
 import {
   currentConversationAtom,
   updateConversationAtom,
+  updateConversationWaitingForResponseAtom,
 } from '@helpers/atoms/Conversation.atom'
 
 export default function useSendChatMessage() {
   const currentConvo = useAtomValue(currentConversationAtom)
   const addNewMessage = useSetAtom(addNewMessageAtom)
   const updateConversation = useSetAtom(updateConversationAtom)
-
+  const updateConvWaiting = useSetAtom(updateConversationWaitingForResponseAtom)
   const [currentPrompt, setCurrentPrompt] = useAtom(currentPromptAtom)
 
   let timeout: any | undefined = undefined
@@ -60,10 +61,15 @@ export default function useSendChatMessage() {
   }
 
   const sendChatMessage = async () => {
+    const convoId = currentConvo?._id
+
+    if (!convoId) return
     setCurrentPrompt('')
+    updateConvWaiting(convoId, true)
+
     const prompt = currentPrompt.trim()
     const newMessage: RawMessage = {
-      conversationId: currentConvo?._id,
+      conversationId: convoId,
       message: prompt,
       user: 'user',
       createdAt: new Date().toISOString(),
@@ -77,10 +83,20 @@ export default function useSendChatMessage() {
     events.emit(EventName.OnNewMessageRequest, newMessage)
 
     if (!currentConvo?.summary && currentConvo) {
-      const updatedConv = {
+      const updatedConv: Conversation = {
         ...currentConvo,
+        lastMessage: prompt,
         summary: `Prompt: ${prompt}`,
       }
+
+      updateConversation(updatedConv)
+      await executeSerial(DataService.UpdateConversation, updatedConv)
+    } else {
+      const updatedConv: Conversation = {
+        ...currentConvo,
+        lastMessage: prompt,
+      }
+
       updateConversation(updatedConv)
       await executeSerial(DataService.UpdateConversation, updatedConv)
     }

diff --git a/web/package.json b/web/package.json
@@ -28,10 +28,10 @@
     "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",
+    "lodash": "^4.17.21",
     "lucide-react": "^0.288.0",
     "marked": "^9.1.2",
     "marked-highlight": "^2.0.6",
@@ -42,6 +42,7 @@
     "react": "18.2.0",
     "react-dom": "18.2.0",
     "react-hook-form": "^7.45.4",
+    "react-intersection-observer": "^9.5.2",
     "sass": "^1.69.4",
     "tailwind-merge": "^1.14.0",
     "tailwindcss": "3.3.3",
@@ -50,6 +51,7 @@
   },
   "devDependencies": {
     "@tailwindcss/forms": "^0.5.4",
+    "@types/lodash": "^4.14.200",
     "@types/node": "20.6.5",
     "@types/uuid": "^9.0.6",
     "encoding": "^0.1.13",