import { useCallback } from 'react' import { DownloadState, HuggingFaceRepoData, Quantization } from '@janhq/core' import { Badge, Button, Progress } from '@janhq/joi' import { useAtomValue, useSetAtom } from 'jotai' import { twMerge } from 'tailwind-merge' import { MainViewState } from '@/constants/screens' import { useCreateNewThread } from '@/hooks/useCreateNewThread' import useDownloadModel from '@/hooks/useDownloadModel' import { modelDownloadStateAtom } from '@/hooks/useDownloadState' import { formatDownloadPercentage, toGibibytes } from '@/utils/converter' import { normalizeModelId } from '@/utils/model' import { mainViewStateAtom } from '@/helpers/atoms/App.atom' import { assistantsAtom } from '@/helpers/atoms/Assistant.atom' import { importHuggingFaceModelStageAtom } from '@/helpers/atoms/HuggingFace.atom' import { downloadedModelsAtom, getDownloadingModelAtom, } from '@/helpers/atoms/Model.atom' type Props = { index: number repoData: HuggingFaceRepoData downloadUrl: string fileName: string fileSize?: number quantization?: Quantization } const ModelDownloadRow: React.FC = ({ downloadUrl, fileName, fileSize = 0, quantization, }) => { const downloadedModels = useAtomValue(downloadedModelsAtom) const { downloadModel, abortModelDownload } = useDownloadModel() const allDownloadStates = useAtomValue(modelDownloadStateAtom) const downloadState: DownloadState | undefined = allDownloadStates[fileName] const downloadingModels = useAtomValue(getDownloadingModelAtom) const { requestCreateNewThread } = useCreateNewThread() const setMainViewState = useSetAtom(mainViewStateAtom) const assistants = useAtomValue(assistantsAtom) const downloadedModel = downloadedModels.find((md) => md.id === fileName) const isDownloading = downloadingModels.some((md) => md === fileName) const setHfImportingStage = useSetAtom(importHuggingFaceModelStageAtom) const onAbortDownloadClick = useCallback(() => { if (downloadUrl) { abortModelDownload(normalizeModelId(downloadUrl)) } }, [downloadUrl, abortModelDownload]) const onDownloadClick = useCallback(async () => { if (downloadUrl) { downloadModel( downloadUrl, normalizeModelId(downloadUrl), normalizeModelId(downloadUrl) ) } }, [downloadUrl, downloadModel]) const onUseModelClick = useCallback(async () => { if (assistants.length === 0) { alert('No assistant available') return } await requestCreateNewThread(assistants[0], downloadedModel) setMainViewState(MainViewState.Thread) setHfImportingStage('NONE') }, [ assistants, downloadedModel, requestCreateNewThread, setMainViewState, setHfImportingStage, ]) if (!downloadUrl) { return null } return (
{quantization && ( {quantization} )}

{fileName}

{toGibibytes(fileSize)}
{downloadedModel ? ( ) : isDownloading ? ( ) : ( )}
) } export default ModelDownloadRow