From fdab8af057f80cf1ccaae0dc42c4e5161925f51e Mon Sep 17 00:00:00 2001
From: Faisal Amir <urmauur@gmail.com>
Date: Mon, 12 Aug 2024 17:49:21 +0700
Subject: [PATCH] feat: allow user configure remote model from my model (#3348)

* feat: allow user configure remote model from my model

* chore: fix linter

* chore: fix linter

* chore: add dedpendecies useCallback model item
---
 web/containers/Layout/TopPanel/index.tsx      |   2 +-
 web/containers/ModelDropdown/ModelSection.tsx |   3 +-
 web/hooks/useGetModelsByEngine.ts             |   2 +-
 .../Settings/MyModels/ModelGroup/index.tsx    | 169 ++++++++++++++++++
 .../Settings/MyModels/ModelItem/index.tsx     |  59 +++++-
 web/screens/Settings/MyModels/index.tsx       |  27 +--
 web/utils/model-engine.ts                     |   6 +
 7 files changed, 240 insertions(+), 28 deletions(-)
 create mode 100644 web/screens/Settings/MyModels/ModelGroup/index.tsx

diff --git a/web/containers/Layout/TopPanel/index.tsx b/web/containers/Layout/TopPanel/index.tsx
index bed3ea604..7cb264f47 100644
--- a/web/containers/Layout/TopPanel/index.tsx
+++ b/web/containers/Layout/TopPanel/index.tsx
@@ -69,7 +69,7 @@ const TopPanel = () => {
   }, [activeThread?.id, setActiveThread, threads])
 
   const onCreateThreadClicked = useCallback(async () => {
-    if (!assistants || assistants.length) {
+    if (!assistants || !assistants.length) {
       toaster({
         title: 'No assistant available.',
         description: `Could not create a new thread. Please add an assistant.`,
diff --git a/web/containers/ModelDropdown/ModelSection.tsx b/web/containers/ModelDropdown/ModelSection.tsx
index 925531a6e..63a2512bf 100644
--- a/web/containers/ModelDropdown/ModelSection.tsx
+++ b/web/containers/ModelDropdown/ModelSection.tsx
@@ -129,7 +129,7 @@ const ModelSection: React.FC<Props> = ({
             {engineName}
           </h6>
         </div>
-        <div className="flex items-center">
+        <div className="flex items-center gap-x-0.5">
           {isRemoteEngine && (
             <Button theme="icon" variant="outline" onClick={onSettingClick}>
               {isEngineReady ? (
@@ -165,7 +165,6 @@ const ModelSection: React.FC<Props> = ({
       <ul>
         {models.map((model) => {
           if (!showModel) return null
-
           return (
             <li
               key={model.model}
diff --git a/web/hooks/useGetModelsByEngine.ts b/web/hooks/useGetModelsByEngine.ts
index 52ef4e9fb..3c2c0962f 100644
--- a/web/hooks/useGetModelsByEngine.ts
+++ b/web/hooks/useGetModelsByEngine.ts
@@ -18,7 +18,7 @@ const useGetModelsByEngine = () => {
 
   // TODO: this function needs to be clean up
   const getModelsByEngine = useCallback(
-    (engine: LlmEngine, searchText: string): Model[] => {
+    (engine: LlmEngine, searchText = ''): Model[] => {
       if (LocalEngines.some((x) => x === engine)) {
         return downloadedModels
           .filter((m) => m.engine === engine)
diff --git a/web/screens/Settings/MyModels/ModelGroup/index.tsx b/web/screens/Settings/MyModels/ModelGroup/index.tsx
new file mode 100644
index 000000000..fe5e9aac6
--- /dev/null
+++ b/web/screens/Settings/MyModels/ModelGroup/index.tsx
@@ -0,0 +1,169 @@
+import React, { useCallback, useEffect, useState } from 'react'
+
+import Image from 'next/image'
+
+import {
+  EngineStatus,
+  LlmEngine,
+  LocalEngine,
+  Model,
+  RemoteEngine,
+  RemoteEngines,
+} from '@janhq/core'
+
+import { Button } from '@janhq/joi'
+import { useAtom, useSetAtom } from 'jotai'
+import {
+  SettingsIcon,
+  ChevronDownIcon,
+  ChevronUpIcon,
+  PlusIcon,
+} from 'lucide-react'
+
+import useEngineQuery from '@/hooks/useEngineQuery'
+import useGetModelsByEngine from '@/hooks/useGetModelsByEngine'
+
+import { getLogoByLocalEngine, getTitleByCategory } from '@/utils/model-engine'
+
+import ModelItem from '../ModelItem'
+
+import { showEngineListModelAtom } from '@/helpers/atoms/Model.atom'
+import { setUpRemoteModelStageAtom } from '@/helpers/atoms/SetupRemoteModel.atom'
+
+type Props = {
+  engine: LlmEngine
+  searchText: string
+}
+
+const ModelGroup: React.FC<Props> = ({ engine, searchText }) => {
+  const [models, setModels] = useState<Model[]>([])
+  const { getModelsByEngine } = useGetModelsByEngine()
+  const setUpRemoteModelStage = useSetAtom(setUpRemoteModelStageAtom)
+  const { data: engineData } = useEngineQuery()
+
+  const [showEngineListModel, setShowEngineListModel] = useAtom(
+    showEngineListModelAtom
+  )
+
+  const engineLogo: string | undefined = models.find(
+    (entry) => entry?.metadata?.logo != null
+  )?.metadata?.logo
+
+  const apiKeyUrl: string | undefined = models.find(
+    (entry) => entry?.metadata?.api_key_url != null
+  )?.metadata?.api_key_url
+
+  const onSettingClick = useCallback(() => {
+    setUpRemoteModelStage('SETUP_API_KEY', engine as unknown as RemoteEngine, {
+      logo: engineLogo,
+      api_key_url: apiKeyUrl,
+    })
+  }, [apiKeyUrl, engine, engineLogo, setUpRemoteModelStage])
+
+  const isEngineReady =
+    engineData?.find((e) => e.name === engine)?.status === EngineStatus.Ready
+
+  const getEngineStatusReady: LlmEngine[] | undefined = engineData
+    ?.filter((e) => e.status === EngineStatus.Ready)
+    .map((x) => x.name as LlmEngine)
+
+  const showModel = showEngineListModel.includes(engine)
+
+  const onClickChevron = useCallback(() => {
+    if (showModel) {
+      setShowEngineListModel((prev) => prev.filter((item) => item !== engine))
+    } else {
+      setShowEngineListModel((prev) => [...prev, engine])
+    }
+  }, [engine, setShowEngineListModel, showModel])
+
+  useEffect(() => {
+    const matchedModels = getModelsByEngine(engine, searchText)
+    setModels(matchedModels)
+    setShowEngineListModel((prev) => [
+      ...prev,
+      ...(getEngineStatusReady as LlmEngine[]),
+    ])
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [getModelsByEngine, engine, searchText, setShowEngineListModel])
+
+  const engineName = getTitleByCategory(engine)
+  const localEngineLogo = getLogoByLocalEngine(engine as LocalEngine)
+  const isRemoteEngine = RemoteEngines.includes(engine as RemoteEngine)
+
+  if (models.length === 0) return null
+
+  return (
+    <div className="w-full py-3">
+      <div className="mb-2 flex justify-between pr-2">
+        <div
+          className="flex cursor-pointer items-center gap-2 pl-3"
+          onClick={onClickChevron}
+        >
+          {!isRemoteEngine && (
+            <Image
+              className="h-6 w-6 flex-shrink-0 rounded-full object-cover"
+              width={48}
+              height={48}
+              src={localEngineLogo as string}
+              alt="logo"
+            />
+          )}
+          {engineLogo && (
+            <Image
+              className="h-6 w-6 flex-shrink-0 rounded-full object-cover"
+              width={48}
+              height={48}
+              src={engineLogo}
+              alt="logo"
+            />
+          )}
+          <h6 className="pr-3 font-medium text-[hsla(var(--text-secondary))]">
+            {engineName}
+          </h6>
+        </div>
+        <div className="flex items-center gap-x-0.5">
+          {isRemoteEngine && (
+            <Button theme="icon" onClick={onSettingClick} variant="outline">
+              {isEngineReady ? (
+                <SettingsIcon
+                  size={14}
+                  className="text-[hsla(var(--text-secondary))]"
+                />
+              ) : (
+                <PlusIcon
+                  size={14}
+                  className="text-[hsla(var(--text-secondary))]"
+                />
+              )}
+            </Button>
+          )}
+          {!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>
+      <div>
+        {models.map((model) => {
+          if (!showModel) return null
+
+          return <ModelItem model={model} key={model.id} />
+        })}
+      </div>
+    </div>
+  )
+}
+
+export default ModelGroup
diff --git a/web/screens/Settings/MyModels/ModelItem/index.tsx b/web/screens/Settings/MyModels/ModelItem/index.tsx
index 4b53f7c55..4c2519e86 100644
--- a/web/screens/Settings/MyModels/ModelItem/index.tsx
+++ b/web/screens/Settings/MyModels/ModelItem/index.tsx
@@ -1,6 +1,12 @@
 import { memo, useCallback, useMemo, useState } from 'react'
 
-import { LocalEngines, Model } from '@janhq/core'
+import {
+  EngineStatus,
+  LocalEngines,
+  Model,
+  RemoteEngine,
+  RemoteEngines,
+} from '@janhq/core'
 import { Badge, Button, useClickOutside } from '@janhq/joi'
 
 import { useAtomValue, useSetAtom } from 'jotai'
@@ -12,12 +18,19 @@ import {
 } from 'lucide-react'
 import { twMerge } from 'tailwind-merge'
 
+import { toaster } from '@/containers/Toast'
+
+import useAssistantQuery from '@/hooks/useAssistantQuery'
+import useEngineQuery from '@/hooks/useEngineQuery'
 import useModelStart from '@/hooks/useModelStart'
 import useModelStop from '@/hooks/useModelStop'
 import useModels from '@/hooks/useModels'
 
+import useThreadCreateMutation from '@/hooks/useThreadCreateMutation'
+
 import { showWarningMultipleModelModalAtom } from '@/screens/HubScreen2/components/WarningMultipleModelModal'
 
+import { MainViewState, mainViewStateAtom } from '@/helpers/atoms/App.atom'
 import { activeModelsAtom } from '@/helpers/atoms/Model.atom'
 
 type Props = {
@@ -33,6 +46,14 @@ const ModelItem: React.FC<Props> = ({ model }) => {
   const stopModel = useModelStop()
   const [more, setMore] = useState(false)
   const { deleteModel } = useModels()
+  const { data: engineData } = useEngineQuery()
+  const createThreadMutation = useThreadCreateMutation()
+  const { data: assistants } = useAssistantQuery()
+  const setMainViewState = useSetAtom(mainViewStateAtom)
+  const isRemoteEngine = RemoteEngines.includes(model.engine as RemoteEngine)
+  const isEngineReady =
+    engineData?.find((e) => e.name === model.engine)?.status ===
+    EngineStatus.Ready
 
   const [menu, setMenu] = useState<HTMLDivElement | null>(null)
   const [toggle, setToggle] = useState<HTMLDivElement | null>(null)
@@ -82,6 +103,33 @@ const ModelItem: React.FC<Props> = ({ model }) => {
     (e) => model.engine != null && e === model.engine
   )
 
+  const onClickCloudModel = useCallback(async () => {
+    if (!isRemoteEngine) return null
+    if (!model || !engineData) return
+    if (!assistants || !assistants.length) {
+      toaster({
+        title: 'No assistant available.',
+        description: `Could not create a new thread. Please add an assistant.`,
+        type: 'error',
+      })
+      return
+    }
+
+    await createThreadMutation.mutateAsync({
+      modelId: model.model,
+      assistant: assistants[0],
+    })
+
+    setMainViewState(MainViewState.Thread)
+  }, [
+    assistants,
+    createThreadMutation,
+    engineData,
+    isRemoteEngine,
+    model,
+    setMainViewState,
+  ])
+
   return (
     <div className="border border-b-0 border-[hsla(var(--app-border))] bg-[hsla(var(--tertiary-bg))] p-4 first:rounded-t-lg last:rounded-b-lg last:border-b">
       <div className="flex flex-col items-start justify-start gap-4 sm:flex-row sm:items-center sm:justify-between">
@@ -89,11 +137,14 @@ const ModelItem: React.FC<Props> = ({ model }) => {
           <div className="flex w-full items-center justify-between">
             <h6
               className={twMerge(
-                'line-clamp-1 max-w-[200px] font-medium',
-                model.engine !== 'cortex.llamacpp' &&
-                  'max-w-none text-[hsla(var(--text-secondary))]'
+                'line-clamp-1 font-medium text-[hsla(var(--text-secondary))]',
+                isLocalModel && 'max-w-[200px]',
+                isRemoteEngine && !isEngineReady
+                  ? 'cursor-not-allowed text-[hsla(var(--text-tertiary))]'
+                  : 'cursor-pointer'
               )}
               title={model.model}
+              onClick={onClickCloudModel}
             >
               {model.model}
             </h6>
diff --git a/web/screens/Settings/MyModels/index.tsx b/web/screens/Settings/MyModels/index.tsx
index 24d657b91..0392fc504 100644
--- a/web/screens/Settings/MyModels/index.tsx
+++ b/web/screens/Settings/MyModels/index.tsx
@@ -18,7 +18,7 @@ import useDropModelBinaries from '@/hooks/useDropModelBinaries'
 
 import { setImportModelStageAtom } from '@/hooks/useImportModel'
 
-import ModelItem from './ModelItem'
+import ModelGroup from './ModelGroup'
 
 import { MainViewState, mainViewStateAtom } from '@/helpers/atoms/App.atom'
 import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom'
@@ -123,27 +123,14 @@ const MyModels = () => {
               )}
             </>
           ) : (
-            <div className="relative w-full">
+            <div className="relative mt-4 w-full">
               {LlmEngines.map((engine) => {
-                const modelByEngine = filteredDownloadedModels.filter(
-                  (x) => x.engine === engine
-                )
-
-                if (modelByEngine.length === 0) return null
-
                 return (
-                  <div className="my-6" key={engine}>
-                    <div className="flex flex-col items-start justify-start gap-2 sm:flex-row sm:items-center sm:justify-between">
-                      <h6 className="text-base font-semibold capitalize">
-                        {engine}
-                      </h6>
-                    </div>
-                    <div className="mt-2">
-                      {modelByEngine.map((model) => (
-                        <ModelItem key={model.model} model={model} />
-                      ))}
-                    </div>
-                  </div>
+                  <ModelGroup
+                    engine={engine}
+                    key={engine}
+                    searchText={searchText}
+                  />
                 )
               })}
             </div>
diff --git a/web/utils/model-engine.ts b/web/utils/model-engine.ts
index 978f34677..9c6630ed0 100644
--- a/web/utils/model-engine.ts
+++ b/web/utils/model-engine.ts
@@ -7,6 +7,12 @@ export const getTitleByCategory = (category: ModelHubCategory) => {
   switch (category) {
     case 'cortex.llamacpp':
       return 'llama.cpp'
+    case 'cortex.onnx':
+      return 'Onnx'
+    case 'cortex.tensorrt-llm':
+      return 'Tensorrt-llm'
+    case 'triton_trtllm':
+      return 'Triton-trtllm'
     case 'BuiltInModels':
       return 'Built-in Models'
     case 'HuggingFace':