fix: stop button is gone after switching to a new model (#4356)

janhq · Dec 30, 2024 · 81ad90b · 81ad90b
1 parent 9b8d8ab
commit 81ad90b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 14 deletions.
diff --git a/web/helpers/atoms/Thread.atom.ts b/web/helpers/atoms/Thread.atom.ts
@@ -204,6 +204,14 @@ export const resetThreadWaitingForResponseAtom = atom(null, (get, set) => {
   set(threadStatesAtom, currentState)
 })
 
+/**
+ * Reset all generating states
+ **/
+export const resetGeneratingResponseAtom = atom(null, (get, set) => {
+  set(resetThreadWaitingForResponseAtom)
+  set(isGeneratingResponseAtom, false)
+})
+
 /**
  * Update the thread last message
  */

diff --git a/web/hooks/useActiveModel.ts b/web/hooks/useActiveModel.ts
@@ -38,10 +38,6 @@ export function useActiveModel() {
   const pendingModelLoad = useRef(false)
   const isVulkanEnabled = useAtomValue(vulkanEnabledAtom)
   const activeAssistant = useAtomValue(activeAssistantAtom)
-  const setGeneratingResponse = useSetAtom(isGeneratingResponseAtom)
-  const resetThreadWaitingForResponseState = useSetAtom(
-    resetThreadWaitingForResponseAtom
-  )
 
   const downloadedModelsRef = useRef<Model[]>([])
 
@@ -147,8 +143,6 @@ export function useActiveModel() {
         return
 
       const engine = EngineManager.instance().get(stoppingModel.engine)
-      setGeneratingResponse(false)
-      resetThreadWaitingForResponseState()
       return engine
         ?.unloadModel(stoppingModel)
         .catch((e) => console.error(e))
@@ -158,14 +152,7 @@ export function useActiveModel() {
           pendingModelLoad.current = false
         })
     },
-    [
-      activeModel,
-      setStateModel,
-      setActiveModel,
-      stateModel,
-      setGeneratingResponse,
-      resetThreadWaitingForResponseState,
-    ]
+    [activeModel, setStateModel, setActiveModel, stateModel]
   )
 
   const stopInference = useCallback(async () => {

diff --git a/web/screens/Thread/ThreadCenterPanel/AssistantSetting/index.tsx b/web/screens/Thread/ThreadCenterPanel/AssistantSetting/index.tsx
@@ -12,6 +12,7 @@ import { activeAssistantAtom } from '@/helpers/atoms/Assistant.atom'
 import {
   activeThreadAtom,
   engineParamsUpdateAtom,
+  resetGeneratingResponseAtom,
 } from '@/helpers/atoms/Thread.atom'
 
 type Props = {
@@ -24,6 +25,7 @@ const AssistantSetting: React.FC<Props> = ({ componentData }) => {
   const { updateThreadMetadata } = useCreateNewThread()
   const { stopModel } = useActiveModel()
   const setEngineParamsUpdate = useSetAtom(engineParamsUpdateAtom)
+  const resetGenerating = useSetAtom(resetGeneratingResponseAtom)
 
   const onValueChanged = useCallback(
     (key: string, value: string | number | boolean | string[]) => {
@@ -32,6 +34,7 @@ const AssistantSetting: React.FC<Props> = ({ componentData }) => {
         componentData.find((x) => x.key === key)?.requireModelReload ?? false
       if (shouldReloadModel) {
         setEngineParamsUpdate(true)
+        resetGenerating()
         stopModel()
       }