jan/web/hooks/useGetModelsByEngine.ts

import { useCallback } from 'react'

import { LlmEngine, LocalEngines, Model } from '@janhq/core'

import { useQueryClient } from '@tanstack/react-query'

import { useAtomValue } from 'jotai'

import { HfModelEntry } from '@/utils/huggingface'

import { cortexHubModelsQueryKey } from './useModelHub'

import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom'

const useGetModelsByEngine = () => {
  const downloadedModels = useAtomValue(downloadedModelsAtom)
  const queryClient = useQueryClient()

  // TODO: this function needs to be clean up
  const getModelsByEngine = useCallback(
    (engine: LlmEngine, searchText = ''): Model[] => {
      if (LocalEngines.some((x) => x === engine)) {
        return downloadedModels
          .filter((m) => m.engine === engine)
          .filter((m) => {
            if (searchText.trim() === '') return true
            return (
              m.model?.toLowerCase().includes(searchText) ||
              m.name?.toLowerCase().includes(searchText)
            )
          })
      }

      const availableModels = downloadedModels.filter(
        (m) => m.engine === engine
      )
      // engine is remote engine
      const data = queryClient.getQueryData(cortexHubModelsQueryKey)
      if (!data) return availableModels

      const modelEntries = data as HfModelEntry[]
      const models: Model[] = [...availableModels]
      for (const entry of modelEntries) {
        const entryModel = entry.model
        if (!entryModel) continue
        if (entry.engine !== engine) continue
        if (models.some((m) => m.model === entryModel.model)) continue
        models.push(entryModel)
      }

      return models.filter((m) => {
        if (searchText.trim() === '') return true
        return (
          m.model?.toLowerCase().includes(searchText) ||
          m.name?.toLowerCase().includes(searchText)
        )
      })
    },
    [queryClient, downloadedModels]
  )

  return { getModelsByEngine }
}

export default useGetModelsByEngine