Merge pull request #3976 from janhq/chore/should-keep-model-running-on-message-error

chore: retrieves the exact model running status upon message error

core/src/browser/extensions/model.ts

@@ -15,7 +15,13 @@ export abstract class ModelExtension extends BaseExtension implements ModelInter
   abstract getModels(): Promise<Model[]>
   abstract pullModel(model: string, id?: string, name?: string): Promise<void>
   abstract cancelModelPull(modelId: string): Promise<void>
-  abstract importModel(model: string, modePath: string, name?: string, optionType?: OptionType): Promise<void>
+  abstract importModel(
+    model: string,
+    modePath: string,
+    name?: string,
+    optionType?: OptionType
+  ): Promise<void>
   abstract updateModel(modelInfo: Partial<Model>): Promise<Model>
   abstract deleteModel(model: string): Promise<void>
+  abstract isModelLoaded(model: string): Promise<boolean>
 }

extensions/model-extension/src/cortex.ts

@@ -9,7 +9,12 @@ interface ICortexAPI {
   getModel(model: string): Promise<Model>
   getModels(): Promise<Model[]>
   pullModel(model: string, id?: string, name?: string): Promise<void>
-  importModel(path: string, modelPath: string, name?: string, option?: string): Promise<void>
+  importModel(
+    path: string,
+    modelPath: string,
+    name?: string,
+    option?: string
+  ): Promise<void>
   deleteModel(model: string): Promise<void>
   updateModel(model: object): Promise<void>
   cancelModelPull(model: string): Promise<void>
@@ -141,6 +146,17 @@ export class CortexAPI implements ICortexAPI {
     )
   }
 
+  /**
+   * Check model status
+   * @param model
+   */
+  async getModelStatus(model: string): Promise<boolean> {
+    return this.queue
+      .add(() => ky.get(`${API_URL}/models/status/${model}`))
+      .then((e) => true)
+      .catch(() => false)
+  }
+
   /**
    * Do health check on cortex.cpp
    * @returns
@@ -215,7 +231,7 @@ export class CortexAPI implements ICortexAPI {
     }
     model.metadata = model.metadata ?? {
       tags: [],
-      size: model.size ?? model.metadata?.size ?? 0
+      size: model.size ?? model.metadata?.size ?? 0,
     }
     return model as Model
   }

extensions/model-extension/src/index.ts

@@ -238,6 +238,14 @@ export default class JanModelExtension extends ModelExtension {
     return this.cortexAPI.importModel(model, modelPath, name, option)
   }
 
+  /**
+   * Check model status
+   * @param model
+   */
+  async isModelLoaded(model: string): Promise<boolean> {
+    return this.cortexAPI.getModelStatus(model)
+  }
+
   /**
    * Handle download state from main app
    */

web/containers/Providers/EventHandler.tsx

@@ -16,6 +16,7 @@ import {
   EngineManager,
   InferenceEngine,
   extractInferenceParams,
+  ModelExtension,
 } from '@janhq/core'
 import { useAtomValue, useSetAtom } from 'jotai'
 import { ulid } from 'ulidx'
@@ -180,9 +181,17 @@ export default function EventHandler({ children }: { children: ReactNode }) {
         }
         return
       } else if (message.status === MessageStatus.Error) {
+        ;(async () => {
+          if (
+            !(await extensionManager
+              .get<ModelExtension>(ExtensionTypeEnum.Model)
+              ?.isModelLoaded(activeModelRef.current?.id as string))
+          ) {
             setActiveModel(undefined)
             setStateModel({ state: 'start', loading: false, model: undefined })
           }
+        })()
+      }
       // Mark the thread as not waiting for response
       updateThreadWaiting(message.thread_id, false)
 

web/hooks/useModels.ts

@@ -35,6 +35,10 @@ const useModels = () => {
       const localModels = (await getModels()).map((e) => ({
         ...e,
         name: ModelManager.instance().models.get(e.id)?.name ?? e.id,
+        settings:
+          ModelManager.instance().models.get(e.id)?.settings ?? e.settings,
+        parameters:
+          ModelManager.instance().models.get(e.id)?.parameters ?? e.parameters,
         metadata:
           ModelManager.instance().models.get(e.id)?.metadata ?? e.metadata,
       }))