81fea5665b
* chore: enhance onboarding screen's models * chore: lint fix * chore: correct lint fix command * chore: fix tests
645 lines
24 KiB
TypeScript
645 lines
24 KiB
TypeScript
import { useState, useMemo, useEffect, useCallback, useRef } from 'react'
|
|
|
|
import Image from 'next/image'
|
|
|
|
import { EngineConfig, InferenceEngine } from '@janhq/core'
|
|
import {
|
|
Badge,
|
|
Button,
|
|
Input,
|
|
ScrollArea,
|
|
Tabs,
|
|
useClickOutside,
|
|
} from '@janhq/joi'
|
|
|
|
import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai'
|
|
|
|
import {
|
|
ChevronDownIcon,
|
|
ChevronUpIcon,
|
|
DownloadCloudIcon,
|
|
XIcon,
|
|
} from 'lucide-react'
|
|
import { twMerge } from 'tailwind-merge'
|
|
|
|
import ProgressCircle from '@/containers/Loader/ProgressCircle'
|
|
|
|
import ModelLabel from '@/containers/ModelLabel'
|
|
|
|
import SetupRemoteModel from '@/containers/SetupRemoteModel'
|
|
|
|
import { useActiveModel } from '@/hooks/useActiveModel'
|
|
|
|
import { useCreateNewThread } from '@/hooks/useCreateNewThread'
|
|
import useDownloadModel from '@/hooks/useDownloadModel'
|
|
import { modelDownloadStateAtom } from '@/hooks/useDownloadState'
|
|
import { useGetEngines } from '@/hooks/useEngineManagement'
|
|
|
|
import { useGetModelSources } from '@/hooks/useModelSource'
|
|
import useRecommendedModel from '@/hooks/useRecommendedModel'
|
|
|
|
import useUpdateModelParameters from '@/hooks/useUpdateModelParameters'
|
|
|
|
import { formatDownloadPercentage, toGigabytes } from '@/utils/converter'
|
|
|
|
import { manualRecommendationModel } from '@/utils/model'
|
|
import { getLogoEngine, getTitleByEngine } from '@/utils/modelEngine'
|
|
|
|
import { extractModelName } from '@/utils/modelSource'
|
|
|
|
import { activeAssistantAtom } from '@/helpers/atoms/Assistant.atom'
|
|
import {
|
|
configuredModelsAtom,
|
|
getDownloadingModelAtom,
|
|
selectedModelAtom,
|
|
showEngineListModelAtom,
|
|
} from '@/helpers/atoms/Model.atom'
|
|
|
|
import {
|
|
activeThreadAtom,
|
|
setThreadModelParamsAtom,
|
|
} from '@/helpers/atoms/Thread.atom'
|
|
|
|
type Props = {
|
|
chatInputMode?: boolean
|
|
strictedThread?: boolean
|
|
disabled?: boolean
|
|
}
|
|
|
|
export const modelDropdownStateAtom = atom(false)
|
|
|
|
const ModelDropdown = ({
|
|
disabled,
|
|
chatInputMode,
|
|
strictedThread = true,
|
|
}: Props) => {
|
|
const { downloadModel } = useDownloadModel()
|
|
const [modelDropdownState, setModelDropdownState] = useAtom(
|
|
modelDropdownStateAtom
|
|
)
|
|
|
|
const [searchFilter, setSearchFilter] = useState('local')
|
|
const [searchText, setSearchText] = useState('')
|
|
const [open, setOpen] = useState<boolean>(modelDropdownState)
|
|
const activeThread = useAtomValue(activeThreadAtom)
|
|
const activeAssistant = useAtomValue(activeAssistantAtom)
|
|
const downloadingModels = useAtomValue(getDownloadingModelAtom)
|
|
const [toggle, setToggle] = useState<HTMLDivElement | null>(null)
|
|
const [selectedModel, setSelectedModel] = useAtom(selectedModelAtom)
|
|
const { recommendedModel, downloadedModels } = useRecommendedModel()
|
|
const { sources } = useGetModelSources()
|
|
const [dropdownOptions, setDropdownOptions] = useState<HTMLDivElement | null>(
|
|
null
|
|
)
|
|
|
|
const { engines } = useGetEngines()
|
|
|
|
const downloadStates = useAtomValue(modelDownloadStateAtom)
|
|
const setThreadModelParams = useSetAtom(setThreadModelParamsAtom)
|
|
const { updateModelParameter } = useUpdateModelParameters()
|
|
const searchInputRef = useRef<HTMLInputElement>(null)
|
|
const configuredModels = useAtomValue(configuredModelsAtom)
|
|
const { stopModel } = useActiveModel()
|
|
|
|
const featuredModels = sources?.filter((x) =>
|
|
manualRecommendationModel.includes(x.id)
|
|
)
|
|
const { updateThreadMetadata } = useCreateNewThread()
|
|
|
|
const engineList = useMemo(
|
|
() =>
|
|
Object.entries(engines ?? {}).flatMap((e) => ({
|
|
name: e[0],
|
|
type: e[1][0]?.type === 'remote' ? 'remote' : 'local',
|
|
engine: e[1][0],
|
|
})),
|
|
[engines]
|
|
)
|
|
|
|
useClickOutside(() => handleChangeStateOpen(false), null, [
|
|
dropdownOptions,
|
|
toggle,
|
|
])
|
|
|
|
const [showEngineListModel, setShowEngineListModel] = useAtom(
|
|
showEngineListModelAtom
|
|
)
|
|
|
|
const handleChangeStateOpen = useCallback(
|
|
(state: boolean) => {
|
|
setOpen(state)
|
|
setModelDropdownState(state)
|
|
},
|
|
[setModelDropdownState]
|
|
)
|
|
|
|
const filteredDownloadedModels = useMemo(
|
|
() =>
|
|
configuredModels
|
|
.concat(
|
|
downloadedModels.filter(
|
|
(e) => !configuredModels.some((x) => x.id === e.id)
|
|
)
|
|
)
|
|
.filter((e) =>
|
|
e.name.toLowerCase().includes(searchText.toLowerCase().trim())
|
|
)
|
|
.filter((e) => {
|
|
if (searchFilter === 'local') {
|
|
return (
|
|
engineList.find((t) => t.engine?.engine === e.engine)?.type ===
|
|
'local'
|
|
)
|
|
}
|
|
return true
|
|
})
|
|
.sort((a, b) => a.name.localeCompare(b.name))
|
|
.sort((a, b) => {
|
|
const aInDownloadedModels = downloadedModels.some(
|
|
(item) => item.id === a.id
|
|
)
|
|
const bInDownloadedModels = downloadedModels.some(
|
|
(item) => item.id === b.id
|
|
)
|
|
if (aInDownloadedModels && !bInDownloadedModels) {
|
|
return -1
|
|
} else if (!aInDownloadedModels && bInDownloadedModels) {
|
|
return 1
|
|
} else {
|
|
return 0
|
|
}
|
|
}),
|
|
[configuredModels, searchText, searchFilter, downloadedModels, engineList]
|
|
)
|
|
|
|
useEffect(() => {
|
|
if (modelDropdownState && chatInputMode) {
|
|
setOpen(modelDropdownState)
|
|
}
|
|
}, [chatInputMode, modelDropdownState])
|
|
|
|
useEffect(() => {
|
|
if (open && searchInputRef.current) {
|
|
searchInputRef.current.focus()
|
|
}
|
|
}, [open])
|
|
|
|
useEffect(() => {
|
|
setShowEngineListModel((prev) => [
|
|
...prev,
|
|
...engineList
|
|
.filter((x) => (x.engine?.api_key?.length ?? 0) > 0)
|
|
.map((e) => e.name),
|
|
])
|
|
}, [setShowEngineListModel, engineList])
|
|
|
|
useEffect(() => {
|
|
if (!activeThread) return
|
|
const modelId = activeAssistant?.model?.id
|
|
|
|
const model = downloadedModels.find((model) => model.id === modelId)
|
|
if (model) {
|
|
if (
|
|
engines?.[model.engine]?.[0]?.type === 'local' ||
|
|
(engines?.[model.engine]?.[0]?.api_key?.length ?? 0) > 0
|
|
)
|
|
setSelectedModel(model)
|
|
} else {
|
|
setSelectedModel(undefined)
|
|
}
|
|
}, [
|
|
recommendedModel,
|
|
activeThread,
|
|
downloadedModels,
|
|
setSelectedModel,
|
|
activeAssistant?.model?.id,
|
|
engines,
|
|
])
|
|
|
|
const isLocalEngine = useCallback(
|
|
(engine?: string) => {
|
|
if (!engine) return false
|
|
return engineList.some((t) => t.name === engine && t.type === 'local')
|
|
},
|
|
[engineList]
|
|
)
|
|
|
|
const onClickModelItem = useCallback(
|
|
async (modelId: string) => {
|
|
if (!activeAssistant) return
|
|
const model = downloadedModels.find((m) => m.id === modelId)
|
|
setSelectedModel(model)
|
|
setOpen(false)
|
|
stopModel()
|
|
|
|
if (activeThread) {
|
|
// Change assistand tools based on model support RAG
|
|
updateThreadMetadata({
|
|
...activeThread,
|
|
assistants: [
|
|
{
|
|
...activeAssistant,
|
|
tools: [
|
|
{
|
|
type: 'retrieval',
|
|
enabled: model?.engine === InferenceEngine.cortex,
|
|
settings: {
|
|
...(activeAssistant.tools &&
|
|
activeAssistant.tools[0]?.settings),
|
|
},
|
|
},
|
|
],
|
|
},
|
|
],
|
|
})
|
|
|
|
const contextLength = model?.settings.ctx_len
|
|
? Math.min(8192, model?.settings.ctx_len ?? 8192)
|
|
: undefined
|
|
const overriddenParameters = {
|
|
ctx_len: contextLength,
|
|
max_tokens: contextLength
|
|
? Math.min(model?.parameters.max_tokens ?? 8192, contextLength)
|
|
: model?.parameters.max_tokens,
|
|
}
|
|
|
|
const modelParams = {
|
|
...model?.parameters,
|
|
...model?.settings,
|
|
...overriddenParameters,
|
|
}
|
|
|
|
// Update model parameter to the thread state
|
|
setThreadModelParams(activeThread.id, modelParams)
|
|
|
|
// Update model parameter to the thread file
|
|
if (model)
|
|
updateModelParameter(activeThread, {
|
|
params: modelParams,
|
|
modelId: model.id,
|
|
engine: model.engine,
|
|
})
|
|
}
|
|
},
|
|
[
|
|
activeAssistant,
|
|
downloadedModels,
|
|
setSelectedModel,
|
|
activeThread,
|
|
updateThreadMetadata,
|
|
setThreadModelParams,
|
|
updateModelParameter,
|
|
stopModel,
|
|
]
|
|
)
|
|
|
|
const isDownloadALocalModel = useMemo(
|
|
() =>
|
|
downloadedModels.some((x) =>
|
|
engineList.some((t) => t.name === x.engine && t.type === 'local')
|
|
),
|
|
[downloadedModels, engineList]
|
|
)
|
|
|
|
if (strictedThread && !activeThread) {
|
|
return null
|
|
}
|
|
|
|
return (
|
|
<div
|
|
className={twMerge('relative', disabled && 'pointer-events-none')}
|
|
data-testid="model-selector"
|
|
>
|
|
<div className="flex [&>div]:w-full" ref={setToggle}>
|
|
{chatInputMode ? (
|
|
<Badge
|
|
data-testid="model-selector-badge"
|
|
theme="secondary"
|
|
variant={open ? 'solid' : 'outline'}
|
|
className={twMerge(
|
|
'inline-block max-w-[200px] cursor-pointer overflow-hidden text-ellipsis whitespace-nowrap',
|
|
open && 'border border-transparent'
|
|
)}
|
|
onClick={() => handleChangeStateOpen(!open)}
|
|
>
|
|
<span
|
|
className={twMerge(
|
|
!selectedModel && 'text-[hsla(var(--text-tertiary))]'
|
|
)}
|
|
>
|
|
{selectedModel?.name || 'Select a model'}
|
|
</span>
|
|
</Badge>
|
|
) : (
|
|
<Input
|
|
value={selectedModel?.name || ''}
|
|
className="cursor-pointer"
|
|
placeholder="Select a model"
|
|
disabled={disabled}
|
|
readOnly
|
|
suffixIcon={
|
|
<ChevronDownIcon
|
|
size={14}
|
|
className={twMerge(open && 'rotate-180')}
|
|
/>
|
|
}
|
|
onClick={() => setOpen(!open)}
|
|
/>
|
|
)}
|
|
</div>
|
|
<div
|
|
className={twMerge(
|
|
'absolute right-0 z-20 mt-2 max-h-80 w-full overflow-hidden rounded-lg border border-[hsla(var(--app-border))] bg-[hsla(var(--app-bg))] shadow-sm',
|
|
open ? 'flex' : 'hidden',
|
|
chatInputMode && 'bottom-8 left-0 w-72'
|
|
)}
|
|
ref={setDropdownOptions}
|
|
>
|
|
<div className="w-full">
|
|
<div className="p-2 pb-0">
|
|
<Tabs
|
|
options={[
|
|
{ name: 'On-device', value: 'local' },
|
|
{ name: 'Cloud', value: 'remote' },
|
|
]}
|
|
tabStyle="segmented"
|
|
value={searchFilter as string}
|
|
onValueChange={(value) => setSearchFilter(value)}
|
|
/>
|
|
</div>
|
|
<div className="relative border-b border-[hsla(var(--app-border))] py-2">
|
|
<Input
|
|
placeholder="Search"
|
|
value={searchText}
|
|
ref={searchInputRef}
|
|
className="rounded-none border-x-0 border-b-0 border-t-0 focus-within:ring-0 "
|
|
onChange={(e) => setSearchText(e.target.value)}
|
|
suffixIcon={
|
|
searchText.length > 0 && (
|
|
<XIcon
|
|
size={16}
|
|
className="cursor-pointer"
|
|
onClick={() => setSearchText('')}
|
|
/>
|
|
)
|
|
}
|
|
/>
|
|
</div>
|
|
<ScrollArea className="h-[calc(100%-90px)] w-full">
|
|
{engineList
|
|
.filter((e) => e.type === searchFilter)
|
|
.filter(
|
|
(e) =>
|
|
e.type === 'remote' ||
|
|
e.name === InferenceEngine.cortex_llamacpp ||
|
|
filteredDownloadedModels.some((e) => e.engine === e.name)
|
|
)
|
|
.map((engine, i) => {
|
|
const isConfigured =
|
|
engine.type === 'local' ||
|
|
((engine.engine as EngineConfig).api_key?.length ?? 0) > 1
|
|
const engineLogo = getLogoEngine(engine.name as InferenceEngine)
|
|
const showModel = showEngineListModel.includes(engine.name)
|
|
const onClickChevron = () => {
|
|
if (showModel) {
|
|
setShowEngineListModel((prev) =>
|
|
prev.filter((item) => item !== engine.name)
|
|
)
|
|
} else {
|
|
setShowEngineListModel((prev) => [...prev, engine.name])
|
|
}
|
|
}
|
|
return (
|
|
<div
|
|
className="relative w-full border-t border-[hsla(var(--app-border))] first:border-t-0"
|
|
key={i}
|
|
>
|
|
<div className="mt-2">
|
|
<div className="flex items-center justify-between px-4">
|
|
<div
|
|
className="flex w-full cursor-pointer items-center gap-2 py-1"
|
|
onClick={onClickChevron}
|
|
>
|
|
{engineLogo && (
|
|
<Image
|
|
className="h-6 w-6 flex-shrink-0"
|
|
width={48}
|
|
height={48}
|
|
src={engineLogo}
|
|
alt="logo"
|
|
/>
|
|
)}
|
|
<h6 className="font-medium capitalize text-[hsla(var(--text-secondary))]">
|
|
{getTitleByEngine(engine.name)}
|
|
</h6>
|
|
</div>
|
|
<div className="-mr-2 flex gap-1">
|
|
{engine.type === 'remote' && (
|
|
<SetupRemoteModel
|
|
engine={engine.name as InferenceEngine}
|
|
isConfigured={
|
|
(engine.engine.api_key?.length ?? 0) > 0
|
|
}
|
|
/>
|
|
)}
|
|
{!showModel ? (
|
|
<Button theme="icon" onClick={onClickChevron}>
|
|
<ChevronDownIcon
|
|
size={14}
|
|
className="text-[hsla(var(--text-secondary))]"
|
|
/>
|
|
</Button>
|
|
) : (
|
|
<Button theme="icon" onClick={onClickChevron}>
|
|
<ChevronUpIcon
|
|
size={14}
|
|
className="text-[hsla(var(--text-secondary))]"
|
|
/>
|
|
</Button>
|
|
)}
|
|
</div>
|
|
</div>
|
|
|
|
{engine.type === 'local' &&
|
|
!isDownloadALocalModel &&
|
|
showModel &&
|
|
!searchText.length && (
|
|
<ul className="pb-2">
|
|
{featuredModels?.map((model) => {
|
|
const isDownloading = downloadingModels.some(
|
|
(md) => md === (model.models[0]?.id ?? model.id)
|
|
)
|
|
return (
|
|
<li
|
|
key={model.id}
|
|
className="flex items-center justify-between gap-4 px-3 py-2 hover:bg-[hsla(var(--dropdown-menu-hover-bg))]"
|
|
>
|
|
<div className="flex items-center gap-2">
|
|
<p
|
|
className="max-w-[200px] overflow-hidden truncate whitespace-nowrap capitalize text-[hsla(var(--text-secondary))]"
|
|
title={model.id}
|
|
>
|
|
{extractModelName(model.id)}
|
|
</p>
|
|
<ModelLabel
|
|
size={model.models[0]?.size}
|
|
compact
|
|
/>
|
|
</div>
|
|
<div className="flex items-center gap-2 text-[hsla(var(--text-tertiary))]">
|
|
<span className="font-medium">
|
|
{toGigabytes(model.models[0]?.size)}
|
|
</span>
|
|
{!isDownloading ? (
|
|
<DownloadCloudIcon
|
|
size={18}
|
|
className="cursor-pointer text-[hsla(var(--app-link))]"
|
|
onClick={() =>
|
|
downloadModel(model.models[0]?.id)
|
|
}
|
|
/>
|
|
) : (
|
|
Object.values(downloadStates)
|
|
.filter(
|
|
(x) =>
|
|
x.modelId ===
|
|
(model.models[0]?.id ?? model.id)
|
|
)
|
|
.map((item) => (
|
|
<ProgressCircle
|
|
key={item.modelId}
|
|
percentage={
|
|
formatDownloadPercentage(
|
|
item?.percent,
|
|
{
|
|
hidePercentage: true,
|
|
}
|
|
) as number
|
|
}
|
|
size={100}
|
|
/>
|
|
))
|
|
)}
|
|
</div>
|
|
</li>
|
|
)
|
|
})}
|
|
</ul>
|
|
)}
|
|
|
|
<ul className="pb-2">
|
|
{filteredDownloadedModels
|
|
.filter(
|
|
(x) =>
|
|
x.engine === engine.name ||
|
|
(x.engine === InferenceEngine.nitro &&
|
|
engine.name === InferenceEngine.cortex_llamacpp)
|
|
)
|
|
.filter((y) => {
|
|
if (isLocalEngine(y.engine) && !searchText.length) {
|
|
return downloadedModels.find((c) => c.id === y.id)
|
|
} else {
|
|
return y
|
|
}
|
|
})
|
|
.map((model) => {
|
|
if (!showModel) return null
|
|
const isDownloading = downloadingModels.some(
|
|
(md) => md === model.id
|
|
)
|
|
const isDownloaded = downloadedModels.some(
|
|
(c) => c.id === model.id
|
|
)
|
|
return (
|
|
<>
|
|
{isDownloaded && (
|
|
<li
|
|
key={model.id}
|
|
className={twMerge(
|
|
'flex items-center justify-between gap-4 px-3 py-2 hover:bg-[hsla(var(--dropdown-menu-hover-bg))]',
|
|
!isConfigured
|
|
? 'cursor-not-allowed text-[hsla(var(--text-tertiary))]'
|
|
: 'text-[hsla(var(--text-primary))]'
|
|
)}
|
|
onClick={() => {
|
|
if (
|
|
!isConfigured &&
|
|
engine.type === 'remote'
|
|
)
|
|
return null
|
|
if (isDownloaded) {
|
|
onClickModelItem(model.id)
|
|
}
|
|
}}
|
|
>
|
|
<div className="flex gap-x-2">
|
|
<p
|
|
className={twMerge(
|
|
'max-w-[200px] overflow-hidden truncate whitespace-nowrap',
|
|
!isDownloaded &&
|
|
'text-[hsla(var(--text-secondary))]'
|
|
)}
|
|
title={model.name}
|
|
>
|
|
{model.name}
|
|
</p>
|
|
<ModelLabel
|
|
size={model.metadata?.size}
|
|
compact
|
|
/>
|
|
</div>
|
|
<div className="flex items-center gap-2 text-[hsla(var(--text-tertiary))]">
|
|
{!isDownloaded && (
|
|
<span className="font-medium">
|
|
{toGigabytes(model.metadata?.size)}
|
|
</span>
|
|
)}
|
|
{!isDownloading && !isDownloaded ? (
|
|
<DownloadCloudIcon
|
|
size={18}
|
|
className="cursor-pointer text-[hsla(var(--app-link))]"
|
|
onClick={() =>
|
|
downloadModel(
|
|
model.sources[0].url,
|
|
model.id
|
|
)
|
|
}
|
|
/>
|
|
) : (
|
|
Object.values(downloadStates)
|
|
.filter((x) => x.modelId === model.id)
|
|
.map((item) => (
|
|
<ProgressCircle
|
|
key={item.modelId}
|
|
percentage={
|
|
formatDownloadPercentage(
|
|
item?.percent,
|
|
{
|
|
hidePercentage: true,
|
|
}
|
|
) as number
|
|
}
|
|
size={100}
|
|
/>
|
|
))
|
|
)}
|
|
</div>
|
|
</li>
|
|
)}
|
|
</>
|
|
)
|
|
})}
|
|
</ul>
|
|
</div>
|
|
</div>
|
|
)
|
|
})}
|
|
</ScrollArea>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
)
|
|
}
|
|
|
|
export default ModelDropdown
|