From 335aeb04bb07628e37fbe1cd335b880c8d1d7d6b Mon Sep 17 00:00:00 2001 From: Louis Date: Tue, 28 Jan 2025 17:01:38 +0700 Subject: [PATCH] feat: allow users to add remote models --- core/src/types/model/modelEntity.ts | 7 +- web/hooks/useEngineManagement.ts | 68 ++++++++ .../Settings/Engines/ModalAddModel.tsx | 154 ++++++++++++++++++ .../Settings/Engines/RemoteEngineSettings.tsx | 5 +- 4 files changed, 231 insertions(+), 3 deletions(-) create mode 100644 web/screens/Settings/Engines/ModalAddModel.tsx diff --git a/core/src/types/model/modelEntity.ts b/core/src/types/model/modelEntity.ts index 7b67a8e942..482dfa1ac9 100644 --- a/core/src/types/model/modelEntity.ts +++ b/core/src/types/model/modelEntity.ts @@ -1,5 +1,3 @@ -import { FileMetadata } from '../file' - /** * Represents the information about a model. * @stored @@ -70,6 +68,11 @@ export type Model = { */ id: string + /** + * The model identifier, modern version of id. + */ + mode?: string + /** * Human-readable name that is used for UI. */ diff --git a/web/hooks/useEngineManagement.ts b/web/hooks/useEngineManagement.ts index 8367ecd20a..10ba606a43 100644 --- a/web/hooks/useEngineManagement.ts +++ b/web/hooks/useEngineManagement.ts @@ -8,6 +8,10 @@ import { EngineConfig, events, EngineEvent, + ModelSource, + ModelSibling, + Model, + ModelEvent, } from '@janhq/core' import { useAtom } from 'jotai' import { atomWithStorage } from 'jotai/utils' @@ -385,3 +389,67 @@ export const uninstallEngine = async ( throw error } } + +/** + * Add a new remote engine model + * @param name + * @param engine + * @returns + */ +export const addRemoteEngineModel = async (name: string, engine: string) => { + const extension = getExtension() + + if (!extension) { + throw new Error('Extension is not available') + } + + try { + // Call the extension's method + const response = await extension.addRemoteModel({ + id: name, + model: name, + engine: engine as InferenceEngine, + } as unknown as Model) + events.emit(ModelEvent.OnModelsUpdate, { fetch: true }) + return response + } catch (error) { + console.error('Failed to install engine variant:', error) + throw error + } +} + +/** + * Remote model sources + * @returns A Promise that resolves to an object of model sources. + */ +export const useGetEngineModelSources = () => { + const { engines } = useGetEngines() + const downloadedModels = useAtomValue(downloadedModelsAtom) + + return { + sources: Object.entries(engines ?? {}) + ?.filter((e) => e?.[1]?.[0]?.type === 'remote') + .map( + ([key, values]) => + ({ + id: key, + models: ( + downloadedModels.filter((e) => e.engine === values[0]?.engine) ?? + [] + ).map( + (e) => + ({ + id: e.id, + size: e.metadata?.size, + }) as unknown as ModelSibling + ), + metadata: { + id: getTitleByEngine(key as InferenceEngine), + description: getDescriptionByEngine(key as InferenceEngine), + apiKey: values[0]?.api_key, + }, + type: 'cloud', + }) as unknown as ModelSource + ), + } +} diff --git a/web/screens/Settings/Engines/ModalAddModel.tsx b/web/screens/Settings/Engines/ModalAddModel.tsx new file mode 100644 index 0000000000..d1ece63788 --- /dev/null +++ b/web/screens/Settings/Engines/ModalAddModel.tsx @@ -0,0 +1,154 @@ +import { memo, ReactNode, useState } from 'react' + +import { useForm } from 'react-hook-form' + +import Image from 'next/image' + +import { zodResolver } from '@hookform/resolvers/zod' + +import { InferenceEngine, Model } from '@janhq/core' + +import { Button, Input, Modal } from '@janhq/joi' +import { PlusIcon } from 'lucide-react' + +import { z } from 'zod' + +import { + addRemoteEngineModel, + useGetEngines, + useGetRemoteModels, +} from '@/hooks/useEngineManagement' + +import { getLogoEngine, getTitleByEngine } from '@/utils/modelEngine' +import { useAtomValue } from 'jotai' +import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom' + +const modelSchema = z.object({ + modelName: z.string().min(1, 'Model name is required'), +}) + +const ModelAddModel = ({ engine }: { engine: string }) => { + const [open, setOpen] = useState(false) + const { mutate: mutateListEngines } = useGetRemoteModels(engine) + const { engines } = useGetEngines() + const models = useAtomValue(downloadedModelsAtom) + const { + register, + handleSubmit, + formState: { errors }, + setError, + } = useForm({ + resolver: zodResolver(modelSchema), + defaultValues: { + modelName: '', + }, + }) + + const onSubmit = async (data: z.infer) => { + if (models.some((e: Model) => e.id === data.modelName)) { + setError('modelName', { + type: 'manual', + message: 'Model already exists', + }) + return + } + await addRemoteEngineModel(data.modelName, engine) + mutateListEngines() + + setOpen(false) + } + + // Helper to render labels with asterisks for required fields + const renderLabel = ( + prefix: ReactNode, + label: string, + isRequired: boolean, + desc?: string + ) => ( + <> + + {prefix} + {label} + +

+ {desc} + {isRequired && *} +

+ + ) + + return ( + +

Add Model

+ + } + fullPage + open={open} + onOpenChange={() => setOpen(!open)} + trigger={ + + } + className="w-[500px]" + content={ +
+
+
+ + + {errors.modelName && ( +

+ {errors.modelName.message} +

+ )} + +
+ +
+ + +
+
+
+ } + /> + ) +} + +export default memo(ModelAddModel) diff --git a/web/screens/Settings/Engines/RemoteEngineSettings.tsx b/web/screens/Settings/Engines/RemoteEngineSettings.tsx index ea2b90b164..83aa07ba44 100644 --- a/web/screens/Settings/Engines/RemoteEngineSettings.tsx +++ b/web/screens/Settings/Engines/RemoteEngineSettings.tsx @@ -27,6 +27,8 @@ import { updateEngine, useGetEngines } from '@/hooks/useEngineManagement' import { getTitleByEngine } from '@/utils/modelEngine' +import ModalAddModel from './ModalAddModel' + import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom' const RemoteEngineSettings = ({ @@ -194,10 +196,11 @@ const RemoteEngineSettings = ({
-
+
Model
+