fix: models load

core/src/browser/extension.ts

@@ -113,7 +113,6 @@ export abstract class BaseExtension implements ExtensionType {
     for (const model of models) {
       ModelManager.instance().register(model)
     }
-    events.emit(ModelEvent.OnModelsUpdate, {})
   }
 
   /**

extensions/inference-cortex-extension/src/index.ts

@@ -215,7 +215,9 @@ export default class JanInferenceCortexExtension extends LocalOAIEngine {
               // Delay for the state update from cortex.cpp
               // Just to be sure
               setTimeout(() => {
-                events.emit(ModelEvent.OnModelsUpdate, {})
+                events.emit(ModelEvent.OnModelsUpdate, {
+                  fetch: true,
+                })
               }, 500)
             }
           })

extensions/model-extension/src/index.ts

@@ -40,11 +40,6 @@ export default class JanModelExtension extends ModelExtension {
   async onLoad() {
     this.registerSettings(SETTINGS)
 
-    // Try get models from cortex.cpp
-    this.getModels().then((models) => {
-      this.registerModels(models)
-    })
-
     // Listen to app download events
     this.handleDesktopEvents()
   }
@@ -163,19 +158,27 @@ export default class JanModelExtension extends ModelExtension {
       (e) => e.engine === InferenceEngine.nitro
     )
 
-    await this.cortexAPI.getModels().then((models) => {
-      const existingIds = models.map((e) => e.id)
-      toImportModels = toImportModels.filter(
-        (e: Model) => !existingIds.includes(e.id) && !e.settings?.vision_model
-      )
-    })
+    /**
+     * Fetch models from cortex.cpp
+     */
+    var fetchedModels = await this.cortexAPI.getModels().catch(() => [])
+
+    // Checking if there are models to import
+    const existingIds = fetchedModels.map((e) => e.id)
+    toImportModels = toImportModels.filter(
+      (e: Model) => !existingIds.includes(e.id) && !e.settings?.vision_model
+    )
+
+    /**
+     * There is no model to import
+     * just return fetched models
+     */
+    if (!toImportModels.length) return fetchedModels
 
     console.log('To import models:', toImportModels.length)
     /**
      * There are models to import
-     * do not return models from cortex.cpp yet
-     * otherwise it will reset the app cache
-     * */
+     */
     if (toImportModels.length > 0) {
       // Import models
       await Promise.all(
@@ -202,8 +205,6 @@ export default class JanModelExtension extends ModelExtension {
           })
         })
       )
-
-      return currentModels
     }
 
     /**

web/containers/Providers/DataLoader.tsx

@@ -29,17 +29,18 @@ const DataLoader: React.FC<Props> = ({ children }) => {
   const setQuickAskEnabled = useSetAtom(quickAskEnabledAtom)
   const setJanDefaultDataFolder = useSetAtom(defaultJanDataFolderAtom)
   const setJanSettingScreen = useSetAtom(janSettingScreenAtom)
+  const { loadDataModel } = useModels()
 
   useThreads()
   useAssistants()
   useGetSystemResources()
   useLoadTheme()
 
-  const { loadDataModel, isUpdated } = useModels()
   useEffect(() => {
-    // Listen for model updates
+    // Load data once
     loadDataModel()
-  }, [isUpdated, loadDataModel])
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
 
   useEffect(() => {
     window.core?.api

web/containers/Providers/EventListener.tsx

@@ -112,8 +112,8 @@ const EventListenerWrapper = ({ children }: PropsWithChildren) => {
         state.downloadState = 'end'
         setDownloadState(state)
         removeDownloadingModel(state.modelId)
+        events.emit(ModelEvent.OnModelsUpdate, { fetch: true })
       }
-      events.emit(ModelEvent.OnModelsUpdate, {})
     },
     [removeDownloadingModel, setDownloadState]
   )

web/containers/Providers/ModelImportListener.tsx

@@ -43,7 +43,7 @@ const ModelImportListener = ({ children }: PropsWithChildren) => {
   const onImportModelSuccess = useCallback(
     (state: ImportingModel) => {
       if (!state.modelId) return
-      events.emit(ModelEvent.OnModelsUpdate, {})
+      events.emit(ModelEvent.OnModelsUpdate, { fetch: true })
       setImportingModelSuccess(state.importId, state.modelId)
     },
     [setImportingModelSuccess]

web/hooks/useModels.ts

@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useRef } from 'react'
+import { useCallback, useEffect } from 'react'
 
 import {
   ExtensionTypeEnum,
@@ -9,7 +9,7 @@ import {
   ModelManager,
 } from '@janhq/core'
 
-import { useSetAtom } from 'jotai'
+import { useSetAtom, useAtom } from 'jotai'
 
 import { useDebouncedCallback } from 'use-debounce'
 
@@ -27,17 +27,11 @@ import {
  * and updates the atoms accordingly.
  */
 const useModels = () => {
-  const setDownloadedModels = useSetAtom(downloadedModelsAtom)
+  const [downloadedModels, setDownloadedModels] = useAtom(downloadedModelsAtom)
   const setExtensionModels = useSetAtom(configuredModelsAtom)
-  const hasFetchedDownloadedModels = useRef(false) // Track whether the function has been executed
-
-  let isUpdated = false
 
   const getData = useCallback(() => {
-    if (hasFetchedDownloadedModels.current) return
-
     const getDownloadedModels = async () => {
-      hasFetchedDownloadedModels.current = true
       const localModels = (await getModels()).map((e) => ({
         ...e,
         name: ModelManager.instance().models.get(e.id)?.name ?? e.id,
@@ -58,6 +52,8 @@ const useModels = () => {
 
       setDownloadedModels(toUpdate)
 
+      let isUpdated = false
+
       toUpdate.forEach((model) => {
         if (!ModelManager.instance().models.has(model.id)) {
           ModelManager.instance().models.set(model.id, model)
@@ -77,30 +73,41 @@ const useModels = () => {
     // Fetch all data
     getExtensionModels()
     getDownloadedModels()
-  }, [])
+  }, [setDownloadedModels, setExtensionModels])
 
   const reloadData = useDebouncedCallback(() => getData(), 300)
 
+  const updateStates = useCallback(() => {
+    const cachedModels = ModelManager.instance().models.values().toArray()
+    const toUpdate = [
+      ...downloadedModels,
+      ...cachedModels.filter(
+        (e: Model) => !downloadedModels.some((g: Model) => g.id === e.id)
+      ),
+    ]
+
+    setDownloadedModels(toUpdate)
+  }, [downloadedModels, setDownloadedModels])
+
   const getModels = async (): Promise<Model[]> =>
     extensionManager
       .get<ModelExtension>(ExtensionTypeEnum.Model)
       ?.getModels() ?? []
 
   useEffect(() => {
-    // Try get data on mount
-    if (isUpdated) {
-      // Listen for model updates
-      events.on(ModelEvent.OnModelsUpdate, async () => reloadData())
-      return () => {
-        // Remove listener on unmount
-        events.off(ModelEvent.OnModelsUpdate, async () => {})
-      }
+    // Listen for model updates
+    events.on(ModelEvent.OnModelsUpdate, async (data: { fetch?: boolean }) => {
+      if (data.fetch) reloadData()
+      else updateStates()
+    })
+    return () => {
+      // Remove listener on unmount
+      events.off(ModelEvent.OnModelsUpdate, async () => {})
     }
-  }, [isUpdated, reloadData])
+  }, [reloadData, updateStates])
 
   return {
     loadDataModel: getData,
-    isUpdated: isUpdated,
   }
 }