import { useMemo } from 'react' import { ExtensionTypeEnum, ModelExtension } from '@janhq/core' import useSWR from 'swr' import { extensionManager } from '@/extension/ExtensionManager' /** * @returns A Promise that resolves to an object of model sources. */ export function useGetModelSources() { const extension = useMemo( () => extensionManager.get(ExtensionTypeEnum.Model), [] ) const { data: sources, error, mutate, } = useSWR( extension ? 'getSources' : null, () => extension?.getSources().then((e) => e.map((m) => ({ ...m, models: m.models.sort((a, b) => a.size - b.size), })) ), { revalidateOnFocus: false, revalidateOnReconnect: true, } ) return { sources, error, mutate } } export const useModelSourcesMutation = () => { const extension = useMemo( () => extensionManager.get(ExtensionTypeEnum.Model), [] ) /** * Add a new model source * @returns A Promise that resolves to intall of engine. */ const addModelSource = async (source: string) => { try { // Call the extension's method return await extension?.addSource(source) } catch (error) { console.error('Failed to install engine variant:', error) throw error } } /** * Delete a new model source * @returns A Promise that resolves to intall of engine. */ const deleteModelSource = async (source: string) => { try { // Call the extension's method return await extension?.deleteSource(source) } catch (error) { console.error('Failed to install engine variant:', error) throw error } } return { addModelSource, deleteModelSource } }