fix: error while importing local model is not shown (#3294)

* fix: error while importing local model is not shown

Signed-off-by: James <namnh0122@gmail.com>

* fix: failed download should not be added to download state (#3297)

Signed-off-by: James <namnh0122@gmail.com>

---------

Signed-off-by: James <namnh0122@gmail.com>

core/src/types/model/index.ts

@@ -1,4 +1,3 @@
 export * from './modelEntity'
-export * from './modelInterface'
 export * from './modelImport'
 export * from './chatCompletion'

core/src/types/model/modelInterface.ts

@@ -1,52 +0,0 @@
-import { GpuSetting } from '../miscellaneous'
-import { Model } from './modelEntity'
-
-/**
- * Model extension for managing models.
- */
-export interface ModelInterface {
-  /**
-   * Downloads a model.
-   * @param model - The model to download.
-   * @param network - Optional object to specify proxy/whether to ignore SSL certificates.
-   * @returns A Promise that resolves when the model has been downloaded.
-   */
-  downloadModel(
-    model: Model,
-    gpuSettings?: GpuSetting,
-    network?: { ignoreSSL?: boolean; proxy?: string }
-  ): Promise<void>
-
-  /**
-   * Cancels the download of a specific model.
-   * @param {string} modelId - The ID of the model to cancel the download for.
-   * @returns {Promise<void>} A promise that resolves when the download has been cancelled.
-   */
-  cancelModelDownload(modelId: string): Promise<void>
-
-  /**
-   * Deletes a model.
-   * @param modelId - The ID of the model to delete.
-   * @returns A Promise that resolves when the model has been deleted.
-   */
-  deleteModel(modelId: string): Promise<void>
-
-  /**
-   * Saves a model.
-   * @param model - The model to save.
-   * @returns A Promise that resolves when the model has been saved.
-   */
-  saveModel(model: Model): Promise<void>
-
-  /**
-   * Gets a list of downloaded models.
-   * @returns A Promise that resolves with an array of downloaded models.
-   */
-  getDownloadedModels(): Promise<Model[]>
-
-  /**
-   * Gets a list of configured models.
-   * @returns A Promise that resolves with an array of configured models.
-   */
-  getConfiguredModels(): Promise<Model[]>
-}

web/hooks/useCortex.ts

@@ -21,6 +21,7 @@ import { UpdateConfigMutationVariables } from './useEngineMutation'
 import { MessageCreateMutationVariables } from './useMessageCreateMutation'
 import { MessageDeleteMutationVariables } from './useMessageDeleteMutation'
 import { MessageUpdateMutationVariables } from './useMessageUpdateMutation'
+import { DownloadModelMutationVariables } from './useModelDownloadMutation'
 
 import { hostAtom } from '@/helpers/atoms/AppConfig.atom'
 
@@ -246,7 +247,8 @@ const useCortex = () => {
   )
 
   const downloadModel = useCallback(
-    async (modelId: string, fileName?: string, persistedModelId?: string) => {
+    async (variables: DownloadModelMutationVariables) => {
+      const { modelId, fileName, persistedModelId } = variables
       const response = await fetch(`${host}/models/${modelId}/pull`, {
         method: 'POST',
         headers: {

web/hooks/useHfModelFetchAndDownload.ts

@@ -3,15 +3,12 @@ import { useCallback } from 'react'
 import { HuggingFaceRepoData } from '@janhq/core'
 import { useQueryClient } from '@tanstack/react-query'
 
-import { useSetAtom } from 'jotai'
-
 import { toaster } from '@/containers/Toast'
 
 import { fetchHuggingFaceRepoData } from '@/utils/huggingface'
 
-import useCortex from './useCortex'
-import { addDownloadModelStateAtom } from './useDownloadState'
 import { fetchHuggingFaceRepoDataQueryKey } from './useHfRepoDataQuery'
+import useModelDownloadMutation from './useModelDownloadMutation'
 
 /**
  * Fetches the data for a Hugging Face model and downloads it.
@@ -21,8 +18,7 @@ import { fetchHuggingFaceRepoDataQueryKey } from './useHfRepoDataQuery'
  * @param modelHandle The model handle to fetch and download. E.g: "NousResearch/Hermes-2-Theta-Llama-3-8B-GGUF"
  */
 const useHfModelFetchAndDownload = () => {
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
+  const downloadModelMutation = useModelDownloadMutation()
-  const { downloadModel } = useCortex()
   const queryClient = useQueryClient()
 
   const fetchData = useCallback(
@@ -90,15 +86,14 @@ const useHfModelFetchAndDownload = () => {
         .concat('_')
         .concat(recommendedModel.rfilename)
 
-      addDownloadState(persistModelId)
+      downloadModelMutation.mutate({
-      await downloadModel(
+        modelId: modelHandle,
-        modelHandle,
+        fileName: recommendedModel.rfilename,
-        recommendedModel.rfilename,
+        persistedModelId: persistModelId,
-        persistModelId
+      })
-      )
     },
 
-    [addDownloadState, downloadModel, fetchData]
+    [fetchData, downloadModelMutation]
   )
 
   return { fetchDataAndDownload }

web/hooks/useModelDownloadMutation.ts

@@ -0,0 +1,49 @@
+import { useMutation } from '@tanstack/react-query'
+
+import { useSetAtom } from 'jotai'
+
+import { toaster } from '@/containers/Toast'
+
+import useCortex from './useCortex'
+import { addDownloadModelStateAtom } from './useDownloadState'
+
+export type DownloadModelMutationVariables = {
+  modelId: string
+  fileName?: string
+  persistedModelId?: string
+}
+
+const useModelDownloadMutation = () => {
+  const { downloadModel } = useCortex()
+  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
+
+  return useMutation({
+    mutationFn: downloadModel,
+
+    onMutate: (variables) => {
+      console.debug('Downloading model', variables)
+    },
+
+    onSuccess: (data, variables) => {
+      console.debug('Download response success', data, variables)
+
+      const { persistedModelId, modelId } = variables
+      if (persistedModelId) {
+        addDownloadState(persistedModelId)
+      } else {
+        addDownloadState(modelId)
+      }
+    },
+
+    onError: (err, variables) => {
+      console.error('Failed to download model', err, variables)
+      toaster({
+        title: 'Failed to download model',
+        description: err.message,
+        type: 'error',
+      })
+    },
+  })
+}
+
+export default useModelDownloadMutation

web/hooks/useSendMessage.ts

@@ -503,6 +503,7 @@ const useSendMessage = () => {
     chatCompletionStreaming,
     setIsGeneratingResponse,
     setShowWarningMultipleModelModal,
+    setChunkCount,
   ])
 
   const sendMessage = useCallback(

web/screens/HubScreen2/components/BuiltInModelCard.tsx

@@ -8,11 +8,8 @@ import { toaster } from '@/containers/Toast'
 
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
-import useCortex from '@/hooks/useCortex'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
-import {
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage } from '@/utils/converter'
@@ -73,9 +70,8 @@ type DownloadContainerProps = {
 const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const { abortDownload } = useAbortDownload()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const setDownloadLocalModelModalStage = useSetAtom(localModelModalStageAtom)
@@ -93,10 +89,9 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, modelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
+  const onDownloadClick = useCallback(() => {
-    addDownloadState(modelId)
+    downloadModelMutation.mutate({ modelId })
-    await downloadModel(modelId)
+  }, [downloadModelMutation, modelId])
-  }, [downloadModel, addDownloadState, modelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/HubScreen2/components/HfListModel.tsx

@@ -9,14 +9,10 @@ import { toaster } from '@/containers/Toast'
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
 
-import useCortex from '@/hooks/useCortex'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
-
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
 
 import useHfRepoDataQuery from '@/hooks/useHfRepoDataQuery'
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
@@ -97,9 +93,8 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   fileName,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const { abortDownload } = useAbortDownload()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const { data: assistants } = useAssistantQuery()
@@ -122,9 +117,12 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   )
 
   const onDownloadClick = useCallback(async () => {
-    addDownloadState(persistModelId)
+    downloadModelMutation.mutate({
-    await downloadModel(modelHandle, fileName, persistModelId)
+      modelId: modelHandle,
-  }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId])
+      fileName: fileName,
+      persistedModelId: persistModelId,
+    })
+  }, [downloadModelMutation, modelHandle, fileName, persistModelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/HubScreen2/components/ListModel.tsx

@@ -9,16 +9,12 @@ import { toaster } from '@/containers/Toast'
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
 
-import useCortex from '@/hooks/useCortex'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
-
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
 
 import useEngineQuery from '@/hooks/useEngineQuery'
 import useHfEngineToBranchesQuery from '@/hooks/useHfEngineToBranchesQuery'
 
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
@@ -150,9 +146,8 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   branch,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
   const { abortDownload } = useAbortDownload()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const setDownloadLocalModelModalStage = useSetAtom(localModelModalStageAtom)
@@ -172,10 +167,9 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, modelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
+  const onDownloadClick = useCallback(() => {
-    addDownloadState(modelId)
+    downloadModelMutation.mutate({ modelId })
-    await downloadModel(modelId)
+  }, [downloadModelMutation, modelId])
-  }, [downloadModel, addDownloadState, modelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/HubScreen2/components/SliderItem.tsx

@@ -10,11 +10,8 @@ import { toaster } from '@/containers/Toast'
 
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
-import useCortex from '@/hooks/useCortex'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
-import {
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
 import { QuickStartModel } from '@/hooks/useModelHub'
 import useThreads from '@/hooks/useThreads'
 
@@ -77,8 +74,7 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   fileName,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const { createThread } = useThreads()
   const { data: assistants } = useAssistantQuery()
@@ -102,10 +98,13 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, persistModelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
+  const onDownloadClick = useCallback(() => {
-    addDownloadState(persistModelId)
+    downloadModelMutation.mutate({
-    await downloadModel(modelHandle, fileName, persistModelId)
+      modelId: modelHandle,
-  }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId])
+      fileName: fileName,
+      persistedModelId: persistModelId,
+    })
+  }, [downloadModelMutation, modelHandle, fileName, persistModelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/Settings/HuggingFaceRepoDetailModal/ModelDownloadRow/index.tsx

@@ -10,12 +10,8 @@ import { toaster } from '@/containers/Toast'
 import useAbortDownload from '@/hooks/useAbortDownload'
 import useAssistantQuery from '@/hooks/useAssistantQuery'
 
-import useCortex from '@/hooks/useCortex'
+import { downloadStateListAtom } from '@/hooks/useDownloadState'
-
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
-import {
-  addDownloadModelStateAtom,
-  downloadStateListAtom,
-} from '@/hooks/useDownloadState'
 import useThreads from '@/hooks/useThreads'
 
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
@@ -69,8 +65,7 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
   modelHandle,
   fileName,
 }) => {
-  const { downloadModel } = useCortex()
+  const downloadModelMutation = useModelDownloadMutation()
-  const addDownloadState = useSetAtom(addDownloadModelStateAtom)
   const setMainViewState = useSetAtom(mainViewStateAtom)
   const setHfImportingStage = useSetAtom(importHuggingFaceModelStageAtom)
   const { createThread } = useThreads()
@@ -93,10 +88,13 @@ const DownloadContainer: React.FC<DownloadContainerProps> = ({
     [downloadedModels, persistModelId]
   )
 
-  const onDownloadClick = useCallback(async () => {
+  const onDownloadClick = useCallback(() => {
-    addDownloadState(persistModelId)
+    downloadModelMutation.mutate({
-    await downloadModel(modelHandle, fileName, persistModelId)
+      modelId: modelHandle,
-  }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId])
+      fileName: fileName,
+      persistedModelId: persistModelId,
+    })
+  }, [downloadModelMutation, modelHandle, fileName, persistModelId])
 
   const onUseModelClick = useCallback(async () => {
     if (!assistants || assistants.length === 0) {

web/screens/Settings/ImportingModelModal/index.tsx

@@ -5,14 +5,14 @@ import { useAtomValue, useSetAtom } from 'jotai'
 
 import { AlertCircle } from 'lucide-react'
 
-import { toaster } from '@/containers/Toast'
-
 import useCortex from '@/hooks/useCortex'
 import {
   getImportModelStageAtom,
   setImportModelStageAtom,
 } from '@/hooks/useImportModel'
 
+import useModelDownloadMutation from '@/hooks/useModelDownloadMutation'
+
 import ImportingModelItem from './ImportingModelItem'
 
 import {
@@ -20,7 +20,8 @@ import {
   setImportingModelErrorAtom,
 } from '@/helpers/atoms/Model.atom'
 
-const ImportingModelModal = () => {
+const ImportingModelModal: React.FC = () => {
+  const downloadModelMutation = useModelDownloadMutation()
   const { downloadModel } = useCortex()
   const setImportModelStage = useSetAtom(setImportModelStageAtom)
   const setImportModelError = useSetAtom(setImportingModelErrorAtom)
@@ -35,19 +36,16 @@ const ImportingModelModal = () => {
     const importModels = async () => {
       for (const model of importingModels) {
         try {
-          await downloadModel(model.path)
+          await downloadModelMutation.mutateAsync({
+            modelId: model.path,
+          })
         } catch (error) {
           let errorMessage = String(error)
           if (error instanceof Error) {
             errorMessage = error.message
           }
 
-          setImportModelError(model.path, errorMessage)
+          setImportModelError(model.importId, errorMessage)
-          toaster({
-            title: 'Import failed',
-            description: errorMessage,
-            type: 'error',
-          })
         }
       }
     }
@@ -68,7 +66,6 @@ const ImportingModelModal = () => {
             ))}
           </div>
 
-          <div>
           <div className="flex flex-row gap-2 rounded-b-lg border-t border-[hsla(var(--app-border))] py-2">
             <AlertCircle
               size={16}
@@ -80,7 +77,6 @@ const ImportingModelModal = () => {
               outputs.
             </p>
           </div>
-          </div>
         </Fragment>
       }
     />