From 26be941e8426462e1e3a28e5b9bf1f834f462f82 Mon Sep 17 00:00:00 2001 From: NamH Date: Wed, 7 Aug 2024 17:55:01 +0700 Subject: [PATCH] fix: error while importing local model is not shown (#3294) * fix: error while importing local model is not shown Signed-off-by: James * fix: failed download should not be added to download state (#3297) Signed-off-by: James --------- Signed-off-by: James --- core/src/types/model/index.ts | 1 - core/src/types/model/modelInterface.ts | 52 ------------------- web/hooks/useCortex.ts | 4 +- web/hooks/useHfModelFetchAndDownload.ts | 21 +++----- web/hooks/useModelDownloadMutation.ts | 49 +++++++++++++++++ web/hooks/useSendMessage.ts | 1 + .../components/BuiltInModelCard.tsx | 17 +++--- .../HubScreen2/components/HfListModel.tsx | 20 ++++--- .../HubScreen2/components/ListModel.tsx | 18 +++---- .../HubScreen2/components/SliderItem.tsx | 21 ++++---- .../ModelDownloadRow/index.tsx | 22 ++++---- .../Settings/ImportingModelModal/index.tsx | 40 +++++++------- 12 files changed, 120 insertions(+), 146 deletions(-) delete mode 100644 core/src/types/model/modelInterface.ts create mode 100644 web/hooks/useModelDownloadMutation.ts diff --git a/core/src/types/model/index.ts b/core/src/types/model/index.ts index e7f711368b..bcfe476d38 100644 --- a/core/src/types/model/index.ts +++ b/core/src/types/model/index.ts @@ -1,4 +1,3 @@ export * from './modelEntity' -export * from './modelInterface' export * from './modelImport' export * from './chatCompletion' diff --git a/core/src/types/model/modelInterface.ts b/core/src/types/model/modelInterface.ts deleted file mode 100644 index 639c7c8d34..0000000000 --- a/core/src/types/model/modelInterface.ts +++ /dev/null @@ -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 - - /** - * Cancels the download of a specific model. - * @param {string} modelId - The ID of the model to cancel the download for. - * @returns {Promise} A promise that resolves when the download has been cancelled. - */ - cancelModelDownload(modelId: string): Promise - - /** - * 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 - - /** - * Saves a model. - * @param model - The model to save. - * @returns A Promise that resolves when the model has been saved. - */ - saveModel(model: Model): Promise - - /** - * Gets a list of downloaded models. - * @returns A Promise that resolves with an array of downloaded models. - */ - getDownloadedModels(): Promise - - /** - * Gets a list of configured models. - * @returns A Promise that resolves with an array of configured models. - */ - getConfiguredModels(): Promise -} diff --git a/web/hooks/useCortex.ts b/web/hooks/useCortex.ts index 6f10ed2452..72992dd138 100644 --- a/web/hooks/useCortex.ts +++ b/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: { diff --git a/web/hooks/useHfModelFetchAndDownload.ts b/web/hooks/useHfModelFetchAndDownload.ts index dcc5472599..76918656dc 100644 --- a/web/hooks/useHfModelFetchAndDownload.ts +++ b/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 { downloadModel } = useCortex() + const downloadModelMutation = useModelDownloadMutation() const queryClient = useQueryClient() const fetchData = useCallback( @@ -90,15 +86,14 @@ const useHfModelFetchAndDownload = () => { .concat('_') .concat(recommendedModel.rfilename) - addDownloadState(persistModelId) - await downloadModel( - modelHandle, - recommendedModel.rfilename, - persistModelId - ) + downloadModelMutation.mutate({ + modelId: modelHandle, + fileName: recommendedModel.rfilename, + persistedModelId: persistModelId, + }) }, - [addDownloadState, downloadModel, fetchData] + [fetchData, downloadModelMutation] ) return { fetchDataAndDownload } diff --git a/web/hooks/useModelDownloadMutation.ts b/web/hooks/useModelDownloadMutation.ts new file mode 100644 index 0000000000..5b75dea85d --- /dev/null +++ b/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 diff --git a/web/hooks/useSendMessage.ts b/web/hooks/useSendMessage.ts index f71438a4da..a06c22c7dd 100644 --- a/web/hooks/useSendMessage.ts +++ b/web/hooks/useSendMessage.ts @@ -503,6 +503,7 @@ const useSendMessage = () => { chatCompletionStreaming, setIsGeneratingResponse, setShowWarningMultipleModelModal, + setChunkCount, ]) const sendMessage = useCallback( diff --git a/web/screens/HubScreen2/components/BuiltInModelCard.tsx b/web/screens/HubScreen2/components/BuiltInModelCard.tsx index f5b35f3818..e96b866014 100644 --- a/web/screens/HubScreen2/components/BuiltInModelCard.tsx +++ b/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 { - addDownloadModelStateAtom, - downloadStateListAtom, -} from '@/hooks/useDownloadState' +import { downloadStateListAtom } from '@/hooks/useDownloadState' +import useModelDownloadMutation from '@/hooks/useModelDownloadMutation' import useThreads from '@/hooks/useThreads' import { formatDownloadPercentage } from '@/utils/converter' @@ -73,9 +70,8 @@ type DownloadContainerProps = { const DownloadContainer: React.FC = ({ 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 = ({ [downloadedModels, modelId] ) - const onDownloadClick = useCallback(async () => { - addDownloadState(modelId) - await downloadModel(modelId) - }, [downloadModel, addDownloadState, modelId]) + const onDownloadClick = useCallback(() => { + downloadModelMutation.mutate({ modelId }) + }, [downloadModelMutation, modelId]) const onUseModelClick = useCallback(async () => { if (!assistants || assistants.length === 0) { diff --git a/web/screens/HubScreen2/components/HfListModel.tsx b/web/screens/HubScreen2/components/HfListModel.tsx index 233e2fe4c1..7a2c1c93bd 100644 --- a/web/screens/HubScreen2/components/HfListModel.tsx +++ b/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 { - addDownloadModelStateAtom, - downloadStateListAtom, -} from '@/hooks/useDownloadState' +import { 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 = ({ 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 = ({ ) const onDownloadClick = useCallback(async () => { - addDownloadState(persistModelId) - await downloadModel(modelHandle, fileName, persistModelId) - }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId]) + downloadModelMutation.mutate({ + modelId: modelHandle, + fileName: fileName, + persistedModelId: persistModelId, + }) + }, [downloadModelMutation, modelHandle, fileName, persistModelId]) const onUseModelClick = useCallback(async () => { if (!assistants || assistants.length === 0) { diff --git a/web/screens/HubScreen2/components/ListModel.tsx b/web/screens/HubScreen2/components/ListModel.tsx index cf7b21b791..46cb951f20 100644 --- a/web/screens/HubScreen2/components/ListModel.tsx +++ b/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 { - addDownloadModelStateAtom, - downloadStateListAtom, -} from '@/hooks/useDownloadState' +import { 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 = ({ 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 = ({ [downloadedModels, modelId] ) - const onDownloadClick = useCallback(async () => { - addDownloadState(modelId) - await downloadModel(modelId) - }, [downloadModel, addDownloadState, modelId]) + const onDownloadClick = useCallback(() => { + downloadModelMutation.mutate({ modelId }) + }, [downloadModelMutation, modelId]) const onUseModelClick = useCallback(async () => { if (!assistants || assistants.length === 0) { diff --git a/web/screens/HubScreen2/components/SliderItem.tsx b/web/screens/HubScreen2/components/SliderItem.tsx index f6fcb5b067..93a1169962 100644 --- a/web/screens/HubScreen2/components/SliderItem.tsx +++ b/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 { - addDownloadModelStateAtom, - downloadStateListAtom, -} from '@/hooks/useDownloadState' +import { downloadStateListAtom } from '@/hooks/useDownloadState' +import useModelDownloadMutation from '@/hooks/useModelDownloadMutation' import { QuickStartModel } from '@/hooks/useModelHub' import useThreads from '@/hooks/useThreads' @@ -77,8 +74,7 @@ const DownloadContainer: React.FC = ({ modelHandle, fileName, }) => { - const { downloadModel } = useCortex() - const addDownloadState = useSetAtom(addDownloadModelStateAtom) + const downloadModelMutation = useModelDownloadMutation() const setMainViewState = useSetAtom(mainViewStateAtom) const { createThread } = useThreads() const { data: assistants } = useAssistantQuery() @@ -102,10 +98,13 @@ const DownloadContainer: React.FC = ({ [downloadedModels, persistModelId] ) - const onDownloadClick = useCallback(async () => { - addDownloadState(persistModelId) - await downloadModel(modelHandle, fileName, persistModelId) - }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId]) + const onDownloadClick = useCallback(() => { + downloadModelMutation.mutate({ + modelId: modelHandle, + fileName: fileName, + persistedModelId: persistModelId, + }) + }, [downloadModelMutation, modelHandle, fileName, persistModelId]) const onUseModelClick = useCallback(async () => { if (!assistants || assistants.length === 0) { diff --git a/web/screens/Settings/HuggingFaceRepoDetailModal/ModelDownloadRow/index.tsx b/web/screens/Settings/HuggingFaceRepoDetailModal/ModelDownloadRow/index.tsx index b05850c007..5f09bf3e32 100644 --- a/web/screens/Settings/HuggingFaceRepoDetailModal/ModelDownloadRow/index.tsx +++ b/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 { - addDownloadModelStateAtom, - downloadStateListAtom, -} from '@/hooks/useDownloadState' +import { downloadStateListAtom } from '@/hooks/useDownloadState' +import useModelDownloadMutation from '@/hooks/useModelDownloadMutation' import useThreads from '@/hooks/useThreads' import { formatDownloadPercentage, toGibibytes } from '@/utils/converter' @@ -69,8 +65,7 @@ const DownloadContainer: React.FC = ({ modelHandle, fileName, }) => { - const { downloadModel } = useCortex() - const addDownloadState = useSetAtom(addDownloadModelStateAtom) + const downloadModelMutation = useModelDownloadMutation() const setMainViewState = useSetAtom(mainViewStateAtom) const setHfImportingStage = useSetAtom(importHuggingFaceModelStageAtom) const { createThread } = useThreads() @@ -93,10 +88,13 @@ const DownloadContainer: React.FC = ({ [downloadedModels, persistModelId] ) - const onDownloadClick = useCallback(async () => { - addDownloadState(persistModelId) - await downloadModel(modelHandle, fileName, persistModelId) - }, [addDownloadState, downloadModel, modelHandle, fileName, persistModelId]) + const onDownloadClick = useCallback(() => { + downloadModelMutation.mutate({ + modelId: modelHandle, + fileName: fileName, + persistedModelId: persistModelId, + }) + }, [downloadModelMutation, modelHandle, fileName, persistModelId]) const onUseModelClick = useCallback(async () => { if (!assistants || assistants.length === 0) { diff --git a/web/screens/Settings/ImportingModelModal/index.tsx b/web/screens/Settings/ImportingModelModal/index.tsx index 0da36e85d9..2581c76219 100644 --- a/web/screens/Settings/ImportingModelModal/index.tsx +++ b/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) - toaster({ - title: 'Import failed', - description: errorMessage, - type: 'error', - }) + setImportModelError(model.importId, errorMessage) } } } @@ -68,18 +66,16 @@ const ImportingModelModal = () => { ))} -
-
- -

- Own your model configurations, use at your own risk. - Misconfigurations may result in lower quality or unexpected - outputs. -

-
+
+ +

+ Own your model configurations, use at your own risk. + Misconfigurations may result in lower quality or unexpected + outputs. +

}