From 947488de7cd6f477a1096c6f750d3b31b5a6bd3f Mon Sep 17 00:00:00 2001 From: Louis Date: Wed, 8 Jan 2025 14:45:13 +0700 Subject: [PATCH] chore: update engines list on changes --- .../engine-management-extension/src/index.ts | 2 + web/containers/ModelDropdown/index.tsx | 4 +- web/containers/Providers/DataLoader.tsx | 3 + web/helpers/atoms/Engines.atom.ts | 7 +++ web/hooks/useEngines.ts | 57 +++++++++++++++++++ web/hooks/useStarterScreen.ts | 6 +- .../Settings/SettingLeftPanel/index.tsx | 6 +- 7 files changed, 78 insertions(+), 7 deletions(-) create mode 100644 web/helpers/atoms/Engines.atom.ts create mode 100644 web/hooks/useEngines.ts diff --git a/extensions/engine-management-extension/src/index.ts b/extensions/engine-management-extension/src/index.ts index 57625eb32c..0d30bf4eac 100644 --- a/extensions/engine-management-extension/src/index.ts +++ b/extensions/engine-management-extension/src/index.ts @@ -13,6 +13,7 @@ import { joinPath, events, ModelEvent, + EngineEvent, } from '@janhq/core' import ky, { HTTPError } from 'ky' import PQueue from 'p-queue' @@ -295,6 +296,7 @@ export default class JSONEngineManagementExtension extends EngineManagementExten await this.addRemoteEngine(data).catch(console.error) }) ) + events.emit(EngineEvent.OnEngineUpdate, {}) DEFAULT_REMOTE_MODELS.forEach(async (data: Model) => { await this.addRemoteModel(data).catch(() => {}) }) diff --git a/web/containers/ModelDropdown/index.tsx b/web/containers/ModelDropdown/index.tsx index 7a12642ff0..ee8d192114 100644 --- a/web/containers/ModelDropdown/index.tsx +++ b/web/containers/ModelDropdown/index.tsx @@ -43,6 +43,7 @@ import { manualRecommendationModel } from '@/utils/model' import { getLogoEngine } from '@/utils/modelEngine' import { activeAssistantAtom } from '@/helpers/atoms/Assistant.atom' +import { installedEnginesAtom } from '@/helpers/atoms/Engines.atom' import { configuredModelsAtom, getDownloadingModelAtom, @@ -85,7 +86,8 @@ const ModelDropdown = ({ const [dropdownOptions, setDropdownOptions] = useState( null ) - const { engines } = useGetEngines() + + const engines = useAtomValue(installedEnginesAtom) const downloadStates = useAtomValue(modelDownloadStateAtom) const setThreadModelParams = useSetAtom(setThreadModelParamsAtom) diff --git a/web/containers/Providers/DataLoader.tsx b/web/containers/Providers/DataLoader.tsx index 4702949961..01093e4b21 100644 --- a/web/containers/Providers/DataLoader.tsx +++ b/web/containers/Providers/DataLoader.tsx @@ -6,6 +6,7 @@ import { AppConfiguration, getUserHomePath } from '@janhq/core' import { useSetAtom } from 'jotai' import useAssistants from '@/hooks/useAssistants' +import useEngines from '@/hooks/useEngines' import useGetSystemResources from '@/hooks/useGetSystemResources' import useModels from '@/hooks/useModels' import useThreads from '@/hooks/useThreads' @@ -25,6 +26,7 @@ const DataLoader: React.FC = () => { const setJanDefaultDataFolder = useSetAtom(defaultJanDataFolderAtom) const setJanSettingScreen = useSetAtom(janSettingScreenAtom) const { getData: loadModels } = useModels() + const { getData: loadEngines } = useEngines() useThreads() useAssistants() @@ -33,6 +35,7 @@ const DataLoader: React.FC = () => { useEffect(() => { // Load data once loadModels() + loadEngines() // eslint-disable-next-line react-hooks/exhaustive-deps }, []) diff --git a/web/helpers/atoms/Engines.atom.ts b/web/helpers/atoms/Engines.atom.ts new file mode 100644 index 0000000000..dad93e18c3 --- /dev/null +++ b/web/helpers/atoms/Engines.atom.ts @@ -0,0 +1,7 @@ +import { EngineConfig, Engines } from '@janhq/core' +import { atom } from 'jotai' + +/** + * Store all of the installed engines including local and remote engines + */ +export const installedEnginesAtom = atom() diff --git a/web/hooks/useEngines.ts b/web/hooks/useEngines.ts new file mode 100644 index 0000000000..a737b98122 --- /dev/null +++ b/web/hooks/useEngines.ts @@ -0,0 +1,57 @@ +import { useCallback, useEffect } from 'react' + +import { + ExtensionTypeEnum, + Model, + ModelEvent, + ModelExtension, + events, + ModelManager, + EngineEvent, + EngineManagementExtension, + Engines, +} from '@janhq/core' + +import { useSetAtom } from 'jotai' + +import { useDebouncedCallback } from 'use-debounce' + +import { extensionManager } from '@/extension' + +import { installedEnginesAtom } from '@/helpers/atoms/Engines.atom' + +/** + * useModels hook - Handles the state of models + * It fetches the downloaded models, configured models and default model from Model Extension + * and updates the atoms accordingly. + */ +const useEngines = () => { + const setInstalledEngines = useSetAtom(installedEnginesAtom) + + const getData = useCallback(() => { + getEngines().then(setInstalledEngines) + }, [setInstalledEngines]) + + const reloadData = useDebouncedCallback(() => getData(), 300) + + const getEngines = async (): Promise => + extensionManager + .get(ExtensionTypeEnum.Engine) + ?.getEngines() + .catch(() => ({}) as Engines) ?? ({} as Engines) + + useEffect(() => { + // Listen for engine updates + events.on(EngineEvent.OnEngineUpdate, reloadData) + return () => { + // Remove listener on unmount + events.off(EngineEvent.OnEngineUpdate, reloadData) + } + }, [reloadData]) + + return { + getData, + } +} + +export default useEngines diff --git a/web/hooks/useStarterScreen.ts b/web/hooks/useStarterScreen.ts index 518821aedd..0c97d48e94 100644 --- a/web/hooks/useStarterScreen.ts +++ b/web/hooks/useStarterScreen.ts @@ -1,19 +1,19 @@ import { useMemo } from 'react' +import { InferenceEngine, EngineConfig } from '@janhq/core' import { useAtomValue } from 'jotai' import { isLocalEngine } from '@/utils/modelEngine' +import { installedEnginesAtom } from '@/helpers/atoms/Engines.atom' import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom' import { threadsAtom } from '@/helpers/atoms/Thread.atom' -import { InferenceEngine, EngineConfig } from '@janhq/core/.' -import { useGetEngines } from './useEngineManagement' export function useStarterScreen() { const downloadedModels = useAtomValue(downloadedModelsAtom) const threads = useAtomValue(threadsAtom) - const { engines } = useGetEngines() + const engines = useAtomValue(installedEnginesAtom) const remoteEngines = engines && diff --git a/web/screens/Settings/SettingLeftPanel/index.tsx b/web/screens/Settings/SettingLeftPanel/index.tsx index a500c80a01..21d954a4ac 100644 --- a/web/screens/Settings/SettingLeftPanel/index.tsx +++ b/web/screens/Settings/SettingLeftPanel/index.tsx @@ -6,13 +6,13 @@ import { useAtomValue } from 'jotai' import LeftPanelContainer from '@/containers/LeftPanelContainer' -import { useGetEngines } from '@/hooks/useEngineManagement' - import { getTitleByEngine, isLocalEngine } from '@/utils/modelEngine' import SettingItem from './SettingItem' import { extensionManager } from '@/extension' +import { installedEnginesAtom } from '@/helpers/atoms/Engines.atom' + import { showSettingActiveLocalEngineAtom, showSettingActiveRemoteEngineAtom, @@ -20,7 +20,7 @@ import { import { janSettingScreenAtom } from '@/helpers/atoms/Setting.atom' const SettingLeftPanel = () => { - const { engines } = useGetEngines() + const engines = useAtomValue(installedEnginesAtom) const settingScreens = useAtomValue(janSettingScreenAtom) const showSettingActiveLocalEngine = useAtomValue(