141 lines
4.3 KiB
TypeScript
141 lines
4.3 KiB
TypeScript
import { DownloadState } from '@janhq/core'
|
|
import { atom } from 'jotai'
|
|
|
|
import { toaster } from '@/containers/Toast'
|
|
|
|
import {
|
|
configuredModelsAtom,
|
|
downloadedModelsAtom,
|
|
removeDownloadingModelAtom,
|
|
} from '@/helpers/atoms/Model.atom'
|
|
|
|
// download states
|
|
export const modelDownloadStateAtom = atom<Record<string, DownloadState>>({})
|
|
|
|
/**
|
|
* Used to set the download state for a particular model.
|
|
*/
|
|
export const setDownloadStateAtom = atom(
|
|
null,
|
|
(get, set, state: DownloadState) => {
|
|
const currentState = { ...get(modelDownloadStateAtom) }
|
|
|
|
if (state.downloadState === 'end') {
|
|
const modelDownloadState = currentState[state.modelId]
|
|
|
|
const updatedChildren: DownloadState[] =
|
|
modelDownloadState.children!.filter(
|
|
(m) => m.fileName !== state.fileName
|
|
)
|
|
updatedChildren.push(state)
|
|
modelDownloadState.children = updatedChildren
|
|
currentState[state.modelId] = modelDownloadState
|
|
|
|
const isAllChildrenDownloadEnd = modelDownloadState.children?.every(
|
|
(m) => m.downloadState === 'end'
|
|
)
|
|
|
|
if (isAllChildrenDownloadEnd) {
|
|
// download successfully
|
|
delete currentState[state.modelId]
|
|
set(removeDownloadingModelAtom, state.modelId)
|
|
|
|
const model = get(configuredModelsAtom).find(
|
|
(e) => e.id === state.modelId
|
|
)
|
|
if (model) set(downloadedModelsAtom, (prev) => [...prev, model])
|
|
toaster({
|
|
title: 'Download Completed',
|
|
description: `Download ${state.modelId} completed`,
|
|
type: 'success',
|
|
})
|
|
}
|
|
} else if (state.downloadState === 'error') {
|
|
// download error
|
|
delete currentState[state.modelId]
|
|
set(removeDownloadingModelAtom, state.modelId)
|
|
if (state.error === 'aborted') {
|
|
toaster({
|
|
title: 'Cancel Download',
|
|
description: `Model ${state.modelId} download cancelled`,
|
|
type: 'warning',
|
|
})
|
|
} else {
|
|
let error = state.error
|
|
if (
|
|
typeof error?.includes === 'function' &&
|
|
state.error?.includes('certificate')
|
|
) {
|
|
error +=
|
|
'. To fix enable "Ignore SSL Certificates" in Advanced settings.'
|
|
}
|
|
toaster({
|
|
title: 'Download Failed',
|
|
description: `Model ${state.modelId} download failed: ${error}`,
|
|
type: 'error',
|
|
})
|
|
}
|
|
} else {
|
|
// download in progress
|
|
if (state.size.total === 0) {
|
|
// this is initial state, just set the state
|
|
currentState[state.modelId] = state
|
|
set(modelDownloadStateAtom, currentState)
|
|
return
|
|
}
|
|
|
|
const modelDownloadState = currentState[state.modelId]
|
|
if (!modelDownloadState) {
|
|
console.debug('setDownloadStateAtom: modelDownloadState not found')
|
|
return
|
|
}
|
|
|
|
// delete the children if the filename is matched and replace the new state
|
|
const updatedChildren: DownloadState[] =
|
|
modelDownloadState.children!.filter(
|
|
(m) => m.fileName !== state.fileName
|
|
)
|
|
|
|
updatedChildren.push(state)
|
|
|
|
// re-calculate the overall progress if we have all the children download data
|
|
const isAnyChildDownloadNotReady = updatedChildren.some(
|
|
(m) => m.size.total === 0
|
|
)
|
|
|
|
modelDownloadState.children = updatedChildren
|
|
|
|
if (isAnyChildDownloadNotReady) {
|
|
// just update the children
|
|
currentState[state.modelId] = modelDownloadState
|
|
set(modelDownloadStateAtom, currentState)
|
|
|
|
return
|
|
}
|
|
|
|
const parentTotalSize = modelDownloadState.size.total
|
|
if (parentTotalSize === 0) {
|
|
// calculate the total size of the parent by sum all children total size
|
|
const totalSize = updatedChildren.reduce(
|
|
(acc, m) => acc + m.size.total,
|
|
0
|
|
)
|
|
|
|
modelDownloadState.size.total = totalSize
|
|
}
|
|
|
|
// calculate the total transferred size by sum all children transferred size
|
|
const transferredSize = updatedChildren.reduce(
|
|
(acc, m) => acc + m.size.transferred,
|
|
0
|
|
)
|
|
modelDownloadState.size.transferred = transferredSize
|
|
modelDownloadState.percent =
|
|
parentTotalSize === 0 ? 0 : transferredSize / parentTotalSize
|
|
currentState[state.modelId] = modelDownloadState
|
|
}
|
|
|
|
set(modelDownloadStateAtom, currentState)
|
|
}
|
|
)
|