From b64749b4bbffdbb8a4807b39af73bccab2a1d15a Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Wed, 14 May 2025 12:08:33 +0700
Subject: [PATCH] fix: distinguish between hub and models search (#4989)

* fix: distinguish between hub and models search

* chore: refresh models hub when going to hub screen
---
 core/src/browser/extensions/model.ts    |  5 ++++
 extensions/model-extension/src/index.ts |  4 +--
 web/containers/ModelSearch/index.tsx    | 36 ++++++++++++++++++++++---
 web/hooks/useEngineManagement.ts        | 19 +++++++++++++
 web/screens/Hub/index.tsx               | 15 +++++++++--
 5 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/core/src/browser/extensions/model.ts b/core/src/browser/extensions/model.ts
index da5ecd62e..238e5999f 100644
--- a/core/src/browser/extensions/model.ts
+++ b/core/src/browser/extensions/model.ts
@@ -40,4 +40,9 @@ export abstract class ModelExtension
    * Delete a model source
    */
   abstract deleteSource(source: string): Promise<void>
+
+  /**
+   * Fetch models hub
+   */
+  abstract fetchModelsHub(): Promise<void>
 }
diff --git a/extensions/model-extension/src/index.ts b/extensions/model-extension/src/index.ts
index 4362ab9a5..fd1e5581d 100644
--- a/extensions/model-extension/src/index.ts
+++ b/extensions/model-extension/src/index.ts
@@ -67,7 +67,7 @@ export default class JanModelExtension extends ModelExtension {
     }
 
     // Sync with cortexsohub
-    this.fetchCortexsoModels()
+    this.fetchModelsHub()
   }
 
   /**
@@ -450,7 +450,7 @@ export default class JanModelExtension extends ModelExtension {
   /**
    * Fetch models from cortex.so
    */
-  private fetchCortexsoModels = async () => {
+  fetchModelsHub = async () => {
     const models = await this.fetchModels()
 
     return this.queue.add(() =>
diff --git a/web/containers/ModelSearch/index.tsx b/web/containers/ModelSearch/index.tsx
index 3513d52e3..223601cb7 100644
--- a/web/containers/ModelSearch/index.tsx
+++ b/web/containers/ModelSearch/index.tsx
@@ -5,12 +5,23 @@ import { SearchIcon } from 'lucide-react'
 
 import { useDebouncedCallback } from 'use-debounce'
 
+import {
+  useGetModelSources,
+  useModelSourcesMutation,
+} from '@/hooks/useModelSource'
+
+import Spinner from '../Loader/Spinner'
+
 type Props = {
+  supportModelImport?: boolean
   onSearchLocal?: (searchText: string) => void
 }
 
-const ModelSearch = ({ onSearchLocal }: Props) => {
+const ModelSearch = ({ supportModelImport, onSearchLocal }: Props) => {
   const [searchText, setSearchText] = useState('')
+  const [isSearching, setSearching] = useState(false)
+  const { mutate } = useGetModelSources()
+  const { addModelSource } = useModelSourcesMutation()
   const inputRef = useRef<HTMLInputElement | null>(null)
   const debounced = useDebouncedCallback(async () => {
     if (searchText.indexOf('/') === -1) {
@@ -20,6 +31,15 @@ const ModelSearch = ({ onSearchLocal }: Props) => {
     }
     // Attempt to search local
     onSearchLocal?.(searchText)
+
+    setSearching(true)
+    // Attempt to search model source
+    if (supportModelImport)
+      addModelSource(searchText)
+        .then(() => mutate())
+        .then(() => onSearchLocal?.(searchText))
+        .catch(console.debug)
+        .finally(() => setSearching(false))
   }, 300)
 
   const onSearchChanged = useCallback(
@@ -50,8 +70,18 @@ const ModelSearch = ({ onSearchLocal }: Props) => {
   return (
     <Input
       ref={inputRef}
-      prefixIcon={<SearchIcon size={16} />}
-      placeholder="Search models..."
+      prefixIcon={
+        isSearching ? (
+          <Spinner size={16} strokeWidth={2} />
+        ) : (
+          <SearchIcon size={16} />
+        )
+      }
+      placeholder={
+        supportModelImport
+          ? 'Search or enter Hugging Face URL'
+          : 'Search models'
+      }
       onChange={onSearchChanged}
       onKeyDown={onKeyDown}
       value={searchText}
diff --git a/web/hooks/useEngineManagement.ts b/web/hooks/useEngineManagement.ts
index 02044ec65..d9eacb592 100644
--- a/web/hooks/useEngineManagement.ts
+++ b/web/hooks/useEngineManagement.ts
@@ -12,6 +12,7 @@ import {
   ModelEvent,
   ModelSource,
   ModelSibling,
+  ModelExtension,
 } from '@janhq/core'
 import { useAtom, useAtomValue } from 'jotai'
 import { atomWithStorage } from 'jotai/utils'
@@ -497,3 +498,21 @@ export const useRefreshModelList = (engine: string) => {
 
   return { refreshingModels, refreshModels }
 }
+
+export const useFetchModelsHub = () => {
+  const extension = useMemo(
+    () => extensionManager.get<ModelExtension>(ExtensionTypeEnum.Model) ?? null,
+    []
+  )
+
+  const { data, error, mutate } = useSWR(
+    extension ? 'fetchModelsHub' : null,
+    () => extension?.fetchModelsHub(),
+    {
+      revalidateOnFocus: false,
+      revalidateOnReconnect: true,
+    }
+  )
+
+  return { modelsHub: data, error, mutate }
+}
diff --git a/web/screens/Hub/index.tsx b/web/screens/Hub/index.tsx
index ff6bf3a77..e0f8936bc 100644
--- a/web/screens/Hub/index.tsx
+++ b/web/screens/Hub/index.tsx
@@ -25,7 +25,10 @@ import { twMerge } from 'tailwind-merge'
 import CenterPanelContainer from '@/containers/CenterPanelContainer'
 import ModelSearch from '@/containers/ModelSearch'
 
-import { useGetEngineModelSources } from '@/hooks/useEngineManagement'
+import {
+  useFetchModelsHub,
+  useGetEngineModelSources,
+} from '@/hooks/useEngineManagement'
 import { setImportModelStageAtom } from '@/hooks/useImportModel'
 
 import {
@@ -85,6 +88,7 @@ const hubCompatibleAtom = atom(false)
 const HubScreen = () => {
   const { sources } = useGetModelSources()
   const { sources: remoteModelSources } = useGetEngineModelSources()
+  const { mutate: fetchModelsHub } = useFetchModelsHub()
   const { addModelSource } = useModelSourcesMutation()
   const [searchValue, setSearchValue] = useState('')
   const [sortSelected, setSortSelected] = useState('newest')
@@ -268,6 +272,10 @@ const HubScreen = () => {
     },
   })
 
+  useEffect(() => {
+    fetchModelsHub()
+  }, [])
+
   return (
     <CenterPanelContainer>
       <m.div
@@ -422,7 +430,10 @@ const HubScreen = () => {
                 <div className="absolute left-1/2 top-1/2 z-10 mx-auto w-4/5 -translate-x-1/2 -translate-y-1/2 rounded-xl sm:w-2/6">
                   <div className="flex flex-col items-center justify-between gap-2 sm:flex-row">
                     <div className="w-full" ref={dropdownRef}>
-                      <ModelSearch onSearchLocal={onSearchUpdate} />
+                      <ModelSearch
+                        onSearchLocal={onSearchUpdate}
+                        supportModelImport
+                      />
                       <div
                         className={twMerge(
                           'invisible absolute mt-2 max-h-[400px] w-full overflow-y-auto rounded-lg border border-[hsla(var(--app-border))] bg-[hsla(var(--app-bg))] shadow-lg',