dify/web/app/components/header/account-setting/model-provider-page/model-selector/model-search.ts

import type { DefaultModel, Model, ModelItem, TypeWithI18N } from '../declarations'
import type { ModelSelectorModelPredicate } from './types'
import Fuse from 'fuse.js'
import { supportFunctionCall } from '@/utils/tool-call'
import { ModelFeatureEnum } from '../declarations'

type ProviderSearchEntry = {
  provider: string
  labels: string[]
  providerKeys: string[]
}

type ModelSearchEntry = {
  provider: string
  model: string
  normalizedLabels: string[]
}

type SearchMatches = {
  providers: Set<string>
  models: Set<string>
}

type ModelSelectorSearchIndex = {
  search: (query: string) => SearchMatches
}

type FilterModelSelectorModelsParams = {
  aiCreditVisibleProviders: Set<string>
  defaultModel?: DefaultModel
  inputValue: string
  installedModelList: Model[]
  modelPredicate?: ModelSelectorModelPredicate
  scopeFeatures: ModelFeatureEnum[]
  searchIndex: ModelSelectorSearchIndex
}

const providerSearchOptions = {
  ignoreDiacritics: true,
  ignoreLocation: true,
  minMatchCharLength: 2,
  shouldSort: false,
  threshold: 0.25,
  keys: [
    { name: 'labels', weight: 2 },
    { name: 'providerKeys', weight: 1 },
  ],
}

const modelSearchOptions = {
  ignoreDiacritics: true,
  shouldSort: false,
  useExtendedSearch: true,
  keys: ['normalizedLabels'],
}

const normalizeModelSearchValue = (value: string) =>
  value
    .toLowerCase()
    .normalize('NFKD')
    .replace(/[^\p{Letter}\p{Number}]+/gu, '')

const looksLikeModelQuery = (value: string) => /\d/.test(value)

const getLabelSearchValues = (label: TypeWithI18N, language: string) => {
  if (label[language] !== undefined) return [label[language]]

  return Array.from(new Set(Object.values(label)))
}

const getProviderKeySearchValues = (provider: string) => {
  const keys = provider.split('/').filter((part) => part && part !== 'langgenius')

  return Array.from(new Set([...keys, ...keys.map(normalizeModelSearchValue)]))
}

const createModelSearchKey = (provider: string, model: string) => `${provider}/${model}`

const modelSupportsScopeFeatures = (modelItem: ModelItem, scopeFeatures: ModelFeatureEnum[]) => {
  if (scopeFeatures.length === 0) return true

  return scopeFeatures.every((feature) => {
    if (feature === ModelFeatureEnum.toolCall) return supportFunctionCall(modelItem.features)

    return modelItem.features?.includes(feature) ?? false
  })
}

export const createModelSelectorSearchIndex = (
  installedModelList: Model[],
  language: string,
): ModelSelectorSearchIndex => {
  const providerEntries = installedModelList.map<ProviderSearchEntry>((model) => {
    return {
      provider: model.provider,
      labels: getLabelSearchValues(model.label, language),
      providerKeys: getProviderKeySearchValues(model.provider),
    }
  })
  const modelEntries = installedModelList.flatMap<ModelSearchEntry>((model) =>
    model.models.map((modelItem) => {
      const labels = getLabelSearchValues(modelItem.label, language)

      return {
        provider: model.provider,
        model: modelItem.model,
        normalizedLabels: Array.from(
          new Set([modelItem.model, ...labels].map(normalizeModelSearchValue)),
        ),
      }
    }),
  )
  const providerFuse = new Fuse(providerEntries, providerSearchOptions)
  const modelFuse = new Fuse(modelEntries, modelSearchOptions)

  return {
    search: (query) => {
      const trimmedQuery = query.trim()

      if (!trimmedQuery) return { providers: new Set(), models: new Set() }

      const normalizedQuery = normalizeModelSearchValue(trimmedQuery)
      const providerMatches = looksLikeModelQuery(trimmedQuery)
        ? new Set<string>()
        : new Set(providerFuse.search(trimmedQuery).map(({ item }) => item.provider))
      const modelMatches = normalizedQuery
        ? new Set(
            modelFuse
              .search(`'${normalizedQuery}`)
              .map(({ item }) => createModelSearchKey(item.provider, item.model)),
          )
        : new Set<string>()

      return {
        providers: providerMatches,
        models: modelMatches,
      }
    },
  }
}

export const filterModelSelectorModels = ({
  aiCreditVisibleProviders,
  defaultModel,
  inputValue,
  installedModelList,
  modelPredicate,
  scopeFeatures,
  searchIndex,
}: FilterModelSelectorModelsParams) => {
  const trimmedInputValue = inputValue.trim()
  const matches = trimmedInputValue
    ? searchIndex.search(trimmedInputValue)
    : { providers: new Set<string>(), models: new Set<string>() }

  const filtered = installedModelList
    .map((model) => {
      const providerMatched = matches.providers.has(model.provider)
      const filteredModels = model.models
        .filter((modelItem) => {
          if (!trimmedInputValue || providerMatched) return true

          return matches.models.has(createModelSearchKey(model.provider, modelItem.model))
        })
        .filter((modelItem) => modelSupportsScopeFeatures(modelItem, scopeFeatures))
        .filter((modelItem) => modelPredicate?.(model, modelItem) ?? true)

      if (
        (trimmedInputValue && filteredModels.length === 0) ||
        (!trimmedInputValue &&
          filteredModels.length === 0 &&
          !aiCreditVisibleProviders.has(model.provider))
      ) {
        return null
      }

      return { ...model, models: filteredModels }
    })
    .filter((model): model is Model => model !== null)

  if (defaultModel?.provider) {
    filtered.sort((a, b) => {
      const aSelected = a.provider === defaultModel.provider ? 0 : 1
      const bSelected = b.provider === defaultModel.provider ? 0 : 1

      return aSelected - bSelected
    })
  }

  return filtered
}