chore: improve models and threads caching (#3744)

* chore: managing and maintaining models and threads in the cache

* test: add tests for hooks

web/containers/Layout/RibbonPanel/index.tsx

@@ -16,14 +16,12 @@ import { mainViewStateAtom, showLeftPanelAtom } from '@/helpers/atoms/App.atom'
 import { editMessageAtom } from '@/helpers/atoms/ChatMessage.atom'
 import { serverEnabledAtom } from '@/helpers/atoms/LocalServer.atom'
 
+import { isDownloadALocalModelAtom } from '@/helpers/atoms/Model.atom'
 import {
   reduceTransparentAtom,
   selectedSettingAtom,
 } from '@/helpers/atoms/Setting.atom'
-import {
-  isDownloadALocalModelAtom,
-  threadsAtom,
-} from '@/helpers/atoms/Thread.atom'
+import { threadsAtom } from '@/helpers/atoms/Thread.atom'
 
 export default function RibbonPanel() {
   const [mainViewState, setMainViewState] = useAtom(mainViewStateAtom)

web/containers/ModelDropdown/index.tsx

@@ -513,7 +513,7 @@ const ModelDropdown = ({
                           const isDownloading = downloadingModels.some(
                             (md) => md.id === model.id
                           )
-                          const isdDownloaded = downloadedModels.some(
+                          const isDownloaded = downloadedModels.some(
                             (c) => c.id === model.id
                           )
                           return (
@@ -528,7 +528,7 @@ const ModelDropdown = ({
                               onClick={() => {
                                 if (!apiKey && !isLocalEngine(model.engine))
                                   return null
-                                if (isdDownloaded) {
+                                if (isDownloaded) {
                                   onClickModelItem(model.id)
                                 }
                               }}
@@ -537,7 +537,7 @@ const ModelDropdown = ({
                                 <p
                                   className={twMerge(
                                     'line-clamp-1',
-                                    !isdDownloaded &&
+                                    !isDownloaded &&
                                       'text-[hsla(var(--text-secondary))]'
                                   )}
                                   title={model.name}
@@ -547,12 +547,12 @@ const ModelDropdown = ({
                                 <ModelLabel metadata={model.metadata} compact />
                               </div>
                               <div className="flex items-center gap-2 text-[hsla(var(--text-tertiary))]">
-                                {!isdDownloaded && (
+                                {!isDownloaded && (
                                   <span className="font-medium">
                                     {toGibibytes(model.metadata.size)}
                                   </span>
                                 )}
-                                {!isDownloading && !isdDownloaded ? (
+                                {!isDownloading && !isDownloaded ? (
                                   <DownloadCloudIcon
                                     size={18}
                                     className="cursor-pointer text-[hsla(var(--app-link))]"

web/helpers/atoms/Extension.atom.test.ts

@@ -0,0 +1,78 @@
+// Extension.atom.test.ts
+
+import { act, renderHook } from '@testing-library/react'
+import * as ExtensionAtoms from './Extension.atom'
+import { useAtom, useAtomValue, useSetAtom } from 'jotai'
+
+describe('Extension.atom.ts', () => {
+  afterEach(() => {
+    jest.clearAllMocks()
+  })
+
+  describe('installingExtensionAtom', () => {
+    it('should initialize as an empty array', () => {
+      const { result } = renderHook(() => useAtomValue(ExtensionAtoms.installingExtensionAtom))
+      expect(result.current).toEqual([])
+    })
+  })
+
+  describe('setInstallingExtensionAtom', () => {
+    it('should add a new installing extension', () => {
+      const { result: setAtom } = renderHook(() => useSetAtom(ExtensionAtoms.setInstallingExtensionAtom))
+      const { result: getAtom } = renderHook(() => useAtomValue(ExtensionAtoms.installingExtensionAtom))
+
+      act(() => {
+        setAtom.current('ext1', { extensionId: 'ext1', percentage: 0 })
+      })
+
+      expect(getAtom.current).toEqual([{ extensionId: 'ext1', percentage: 0 }])
+    })
+
+    it('should update an existing installing extension', () => {
+      const { result: setAtom } = renderHook(() => useSetAtom(ExtensionAtoms.setInstallingExtensionAtom))
+      const { result: getAtom } = renderHook(() => useAtomValue(ExtensionAtoms.installingExtensionAtom))
+
+      act(() => {
+        setAtom.current('ext1', { extensionId: 'ext1', percentage: 0 })
+        setAtom.current('ext1', { extensionId: 'ext1', percentage: 50 })
+      })
+
+      expect(getAtom.current).toEqual([{ extensionId: 'ext1', percentage: 50 }])
+    })
+  })
+
+  describe('removeInstallingExtensionAtom', () => {
+    it('should remove an installing extension', () => {
+      const { result: setAtom } = renderHook(() => useSetAtom(ExtensionAtoms.setInstallingExtensionAtom))
+      const { result: removeAtom } = renderHook(() => useSetAtom(ExtensionAtoms.removeInstallingExtensionAtom))
+      const { result: getAtom } = renderHook(() => useAtomValue(ExtensionAtoms.installingExtensionAtom))
+
+      act(() => {
+        setAtom.current('ext1', { extensionId: 'ext1', percentage: 0 })
+        setAtom.current('ext2', { extensionId: 'ext2', percentage: 50 })
+        removeAtom.current('ext1')
+      })
+
+      expect(getAtom.current).toEqual([{ extensionId: 'ext2', percentage: 50 }])
+    })
+  })
+
+  describe('inActiveEngineProviderAtom', () => {
+    it('should initialize as an empty array', () => {
+      const { result } = renderHook(() => useAtomValue(ExtensionAtoms.inActiveEngineProviderAtom))
+      expect(result.current).toEqual([])
+    })
+
+    it('should persist value in storage', () => {
+      const { result } = renderHook(() => useAtom(ExtensionAtoms.inActiveEngineProviderAtom))
+
+      act(() => {
+        result.current[1](['provider1', 'provider2'])
+      })
+
+      // Simulate a re-render to check if the value persists
+      const { result: newResult } = renderHook(() => useAtomValue(ExtensionAtoms.inActiveEngineProviderAtom))
+      expect(newResult.current).toEqual(['provider1', 'provider2'])
+    })
+  })
+})

web/helpers/atoms/Model.atom.test.ts

@@ -1,4 +1,4 @@
-import { act, renderHook, waitFor } from '@testing-library/react'
+import { act, renderHook } from '@testing-library/react'
 import * as ModelAtoms from './Model.atom'
 import { useAtom, useAtomValue, useSetAtom } from 'jotai'
 
@@ -24,11 +24,6 @@ describe('Model.atom.ts', () => {
       })
     })
   })
-  describe('activeAssistantModelAtom', () => {
-    it('should initialize as undefined', () => {
-      expect(ModelAtoms.activeAssistantModelAtom.init).toBeUndefined()
-    })
-  })
 
   describe('selectedModelAtom', () => {
     it('should initialize as undefined', () => {

web/helpers/atoms/Model.atom.ts

@@ -1,8 +1,59 @@
 import { ImportingModel, InferenceEngine, Model, ModelFile } from '@janhq/core'
 import { atom } from 'jotai'
+import { atomWithStorage } from 'jotai/utils'
+
+/**
+ * Enum for the keys used to store models in the local storage.
+ */
+enum ModelStorageAtomKeys {
+  DownloadedModels = 'downloadedModels',
+  AvailableModels = 'availableModels',
+}
+//// Models Atom
+/**
+ * Downloaded Models Atom
+ * This atom stores the list of models that have been downloaded.
+ */
+export const downloadedModelsAtom = atomWithStorage<ModelFile[]>(
+  ModelStorageAtomKeys.DownloadedModels,
+  []
+)
+
+/**
+ * Configured Models Atom
+ * This atom stores the list of models that have been configured and available to download
+ */
+export const configuredModelsAtom = atomWithStorage<ModelFile[]>(
+  ModelStorageAtomKeys.AvailableModels,
+  []
+)
+
+export const removeDownloadedModelAtom = atom(
+  null,
+  (get, set, modelId: string) => {
+    const downloadedModels = get(downloadedModelsAtom)
+
+    set(
+      downloadedModelsAtom,
+      downloadedModels.filter((e) => e.id !== modelId)
+    )
+  }
+)
+
+/**
+ * Atom to store the selected model (from ModelDropdown)
+ */
+export const selectedModelAtom = atom<ModelFile | undefined>(undefined)
+
+/**
+ * Atom to store the expanded engine sections (from ModelDropdown)
+ */
+export const showEngineListModelAtom = atom<string[]>([InferenceEngine.nitro])
+
+/// End Models Atom
+/// Model Download Atom
 
 export const stateModel = atom({ state: 'start', loading: false, model: '' })
-export const activeAssistantModelAtom = atom<Model | undefined>(undefined)
 
 /**
  * Stores the list of models which are being downloaded.
@@ -30,28 +81,20 @@ export const removeDownloadingModelAtom = atom(
   }
 )
 
-export const downloadedModelsAtom = atom<ModelFile[]>([])
-
-export const removeDownloadedModelAtom = atom(
-  null,
-  (get, set, modelId: string) => {
-    const downloadedModels = get(downloadedModelsAtom)
-
-    set(
-      downloadedModelsAtom,
-      downloadedModels.filter((e) => e.id !== modelId)
-    )
-  }
-)
-
-export const configuredModelsAtom = atom<ModelFile[]>([])
-
-export const defaultModelAtom = atom<Model | undefined>(undefined)
+/// End Model Download Atom
+/// Model Import Atom
 
 /// TODO: move this part to another atom
 // store the paths of the models that are being imported
 export const importingModelsAtom = atom<ImportingModel[]>([])
 
+// DEPRECATED: Remove when moving to cortex.cpp
+// Default model template when importing
+export const defaultModelAtom = atom<Model | undefined>(undefined)
+
+/**
+ * Importing progress Atom
+ */
 export const updateImportingModelProgressAtom = atom(
   null,
   (get, set, importId: string, percentage: number) => {
@@ -69,6 +112,9 @@ export const updateImportingModelProgressAtom = atom(
   }
 )
 
+/**
+ * Importing error Atom
+ */
 export const setImportingModelErrorAtom = atom(
   null,
   (get, set, importId: string, error: string) => {
@@ -87,6 +133,9 @@ export const setImportingModelErrorAtom = atom(
   }
 )
 
+/**
+ * Importing success Atom
+ */
 export const setImportingModelSuccessAtom = atom(
   null,
   (get, set, importId: string, modelId: string) => {
@@ -105,6 +154,9 @@ export const setImportingModelSuccessAtom = atom(
   }
 )
 
+/**
+ * Update importing model metadata Atom
+ */
 export const updateImportingModelAtom = atom(
   null,
   (
@@ -131,6 +183,9 @@ export const updateImportingModelAtom = atom(
   }
 )
 
-export const selectedModelAtom = atom<ModelFile | undefined>(undefined)
+/// End Model Import Atom
 
-export const showEngineListModelAtom = atom<string[]>([InferenceEngine.nitro])
+/// ModelDropdown States Atom
+export const isDownloadALocalModelAtom = atom<boolean>(false)
+export const isAnyRemoteModelConfiguredAtom = atom<boolean>(false)
+/// End ModelDropdown States Atom