Skip to content

Commit

Permalink
feat: allow users to add remote models
Browse files Browse the repository at this point in the history
  • Loading branch information
louis-menlo committed Jan 28, 2025
1 parent 65c9bb9 commit 335aeb0
Show file tree
Hide file tree
Showing 4 changed files with 231 additions and 3 deletions.
7 changes: 5 additions & 2 deletions core/src/types/model/modelEntity.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import { FileMetadata } from '../file'

/**
* Represents the information about a model.
* @stored
Expand Down Expand Up @@ -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.
*/
Expand Down
68 changes: 68 additions & 0 deletions web/hooks/useEngineManagement.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,10 @@ import {
EngineConfig,
events,
EngineEvent,
ModelSource,
ModelSibling,
Model,
ModelEvent,
} from '@janhq/core'
import { useAtom } from 'jotai'
import { atomWithStorage } from 'jotai/utils'
Expand Down Expand Up @@ -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
),
}
}
154 changes: 154 additions & 0 deletions web/screens/Settings/Engines/ModalAddModel.tsx
Original file line number Diff line number Diff line change
@@ -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'

Check failure on line 22 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-macos

There should be at least one empty line between import groups

Check failure on line 22 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-windows-pr

There should be at least one empty line between import groups

Check failure on line 22 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-ubuntu

There should be at least one empty line between import groups

Check failure on line 22 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / coverage-check

There should be at least one empty line between import groups
import { useAtomValue } from 'jotai'

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-macos

There should be at least one empty line between import groups

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-macos

`jotai` import should occur before import of `lucide-react`

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-windows-pr

There should be at least one empty line between import groups

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-windows-pr

`jotai` import should occur before import of `lucide-react`

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-ubuntu

There should be at least one empty line between import groups

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / test-on-ubuntu

`jotai` import should occur before import of `lucide-react`

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / coverage-check

There should be at least one empty line between import groups

Check failure on line 23 in web/screens/Settings/Engines/ModalAddModel.tsx

View workflow job for this annotation

GitHub Actions / coverage-check

`jotai` import should occur before import of `lucide-react`
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<typeof modelSchema>) => {
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
) => (
<>
<span className="flex flex-row items-center gap-1">
{prefix}
{label}
</span>
<p className="mt-1 font-normal text-[hsla(var(--text-secondary))]">
{desc}
{isRequired && <span className="text-red-500">*</span>}
</p>
</>
)

return (
<Modal
title={
<div>
<p>Add Model</p>
</div>
}
fullPage
open={open}
onOpenChange={() => setOpen(!open)}
trigger={
<Button>
<PlusIcon className="mr-2" size={14} />
Add Model
</Button>
}
className="w-[500px]"
content={
<div>
<form className="mt-8 space-y-6" onSubmit={handleSubmit(onSubmit)}>
<div className="space-y-2">
<label htmlFor="modelName" className="font-semibold">
{renderLabel(
getLogoEngine(engine as InferenceEngine) ? (
<Image
src={getLogoEngine(engine as InferenceEngine) ?? ''}
width={40}
height={40}
alt="Engine logo"
className="h-5 w-5 flex-shrink-0"
/>
) : (
<></>
),
getTitleByEngine(engine as InferenceEngine) ?? engine,
false,
'Model ID'
)}
</label>
<Input placeholder="Enter model ID" {...register('modelName')} />
{errors.modelName && (
<p className="text-sm text-red-500">
{errors.modelName.message}
</p>
)}
<div className="pt-4">
<a
target="_blank"
href={engines?.[engine as InferenceEngine]?.[0]?.url}
className="text-[hsla(var(--app-link))]"
>
See model list from{' '}
{getTitleByEngine(engine as InferenceEngine)}
</a>
</div>
</div>

<div className="mt-8 flex justify-end gap-x-2">
<Button
theme="ghost"
variant="outline"
onClick={() => setOpen(false)}
>
Cancel
</Button>
<Button type="submit">Add</Button>
</div>
</form>
</div>
}
/>
)
}

export default memo(ModelAddModel)
5 changes: 4 additions & 1 deletion web/screens/Settings/Engines/RemoteEngineSettings.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 = ({
Expand Down Expand Up @@ -194,10 +196,11 @@ const RemoteEngineSettings = ({
<div className="mb-3 mt-4 pb-4">
<div className="flex w-full flex-col items-start justify-between sm:flex-row">
<div className="w-full flex-shrink-0 ">
<div className="flex items-center justify-between gap-x-2">
<div className="mb-4 flex items-center justify-between gap-x-2">
<div>
<h6 className="mb-2 line-clamp-1 font-semibold">Model</h6>
</div>
<ModalAddModel engine={name} />
</div>

<div>
Expand Down

0 comments on commit 335aeb0

Please sign in to comment.