From ca0c5f571f011046bcb05fd2e013103a8eb78cfe Mon Sep 17 00:00:00 2001
From: Louis <133622055+louis-jan@users.noreply.github.com>
Date: Wed, 25 Oct 2023 15:51:03 +0700
Subject: [PATCH] fix: model is started but the indicator is not stopped
 loading (#446)

* fix: model is started but the indicator is not stopped loading

* chore: show loading indicator on the loading model row only

* Update README.md

---------

Co-authored-by: hiento09 <136591877+hiento09@users.noreply.github.com>
---
 plugins/data-plugin/README.md                   |  2 +-
 web/app/_components/ModelActionButton/index.tsx | 15 +++++----------
 web/app/_components/ModelRow/index.tsx          |  6 ++++--
 web/hooks/useStartStopModel.ts                  |  8 +++-----
 4 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/plugins/data-plugin/README.md b/plugins/data-plugin/README.md
index 42caeb1aa..2197d9ad2 100644
--- a/plugins/data-plugin/README.md
+++ b/plugins/data-plugin/README.md
@@ -5,4 +5,4 @@
 - index.ts: Main entry point for the plugin.
 - module.ts: Defines the plugin module which would be executed by the main node process.
 - package.json: Defines the plugin metadata.
-- tsconfig.json: Defines the typescript configuration.
\ No newline at end of file
+- tsconfig.json: Defines the typescript configuration.
diff --git a/web/app/_components/ModelActionButton/index.tsx b/web/app/_components/ModelActionButton/index.tsx
index c1a7c64e0..346d0afc2 100644
--- a/web/app/_components/ModelActionButton/index.tsx
+++ b/web/app/_components/ModelActionButton/index.tsx
@@ -1,9 +1,6 @@
-import React, { useState } from 'react'
+import React from 'react'
 import { Button } from '@uikit'
 import ModelActionMenu from '../ModelActionMenu'
-import { useAtomValue } from 'jotai'
-
-import { stateModel } from '@helpers/atoms/Model.atom'
 
 export enum ModelActionType {
   Start = 'Start',
@@ -25,26 +22,24 @@ const modelActionMapper: Record<ModelActionType, ModelActionStyle> = {
 
 type Props = {
   disabled?: boolean
+  loading?: boolean
   type: ModelActionType
   onActionClick: (type: ModelActionType) => void
   onDeleteClick: () => void
 }
 
 const ModelActionButton: React.FC<Props> = ({
-  disabled = false,
+  disabled,
+  loading,
   type,
   onActionClick,
   onDeleteClick,
 }) => {
   const styles = modelActionMapper[type]
-  // const { startingModel } = useStartStopModel()
 
   const onClick = () => {
     onActionClick(type)
   }
-
-  const state = useAtomValue(stateModel)
-
   return (
     <td className="whitespace-nowrap px-3 py-2 text-right">
       <div className="flex items-center justify-end gap-x-4">
@@ -54,7 +49,7 @@ const ModelActionButton: React.FC<Props> = ({
           size="sm"
           themes={styles.title === 'Start' ? 'accent' : 'default'}
           onClick={() => onClick()}
-          loading={state.loading}
+          loading={loading}
         >
           {styles.title} Model
         </Button>
diff --git a/web/app/_components/ModelRow/index.tsx b/web/app/_components/ModelRow/index.tsx
index f0bd4a283..fd7479982 100644
--- a/web/app/_components/ModelRow/index.tsx
+++ b/web/app/_components/ModelRow/index.tsx
@@ -4,7 +4,7 @@ import { useAtomValue } from 'jotai'
 import ModelActionButton, { ModelActionType } from '../ModelActionButton'
 import useStartStopModel from '@hooks/useStartStopModel'
 import useDeleteModel from '@hooks/useDeleteModel'
-import { activeAssistantModelAtom } from '@helpers/atoms/Model.atom'
+import { activeAssistantModelAtom, stateModel } from '@helpers/atoms/Model.atom'
 import { toGigabytes } from '@utils/converter'
 
 type Props = {
@@ -12,9 +12,10 @@ type Props = {
 }
 
 const ModelRow: React.FC<Props> = ({ model }) => {
-  const { loading, startModel, stopModel } = useStartStopModel()
+  const { startModel, stopModel } = useStartStopModel()
   const activeModel = useAtomValue(activeAssistantModelAtom)
   const { deleteModel } = useDeleteModel()
+  const { loading, model: currentModelState } = useAtomValue(stateModel)
 
   let status = ModelStatus.Installed
   if (activeModel && activeModel._id === model._id) {
@@ -58,6 +59,7 @@ const ModelRow: React.FC<Props> = ({ model }) => {
       </td>
       <ModelActionButton
         disabled={loading}
+        loading={currentModelState === model._id ? loading : false}
         type={actionButtonType}
         onActionClick={onModelActionClick}
         onDeleteClick={onDeleteClick}
diff --git a/web/hooks/useStartStopModel.ts b/web/hooks/useStartStopModel.ts
index b6aca7726..743d8ca6d 100644
--- a/web/hooks/useStartStopModel.ts
+++ b/web/hooks/useStartStopModel.ts
@@ -6,7 +6,6 @@ import { useState } from 'react'
 
 export default function useStartStopModel() {
   const [activeModel, setActiveModel] = useAtom(activeAssistantModelAtom)
-  const [loading, setLoading] = useState<boolean>(false)
   const setStateModel = useSetAtom(stateModel)
 
   const startModel = async (modelId: string) => {
@@ -24,12 +23,11 @@ export default function useStartStopModel() {
     if (!model) {
       alert(`Model ${modelId} not found! Please re-download the model first.`)
       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)
@@ -40,7 +38,7 @@ export default function useStartStopModel() {
       )
       setActiveModel(model)
     }
-    setLoading(false)
+    setStateModel((prev) => ({ ...prev, loading: false }))
   }
 
   const stopModel = async (modelId: string) => {
@@ -52,5 +50,5 @@ export default function useStartStopModel() {
     }, 500)
   }
 
-  return { loading, startModel, stopModel }
+  return { startModel, stopModel }
 }