cherry-studio/src/renderer/src/pages/home/Inputbar/tools/components/WebSearchQuickPanelManager.tsx

import { BaiduOutlined, GoogleOutlined } from '@ant-design/icons'
import { loggerService } from '@logger'
import { BingLogo, BochaLogo, ExaLogo, SearXNGLogo, TavilyLogo, ZhipuLogo } from '@renderer/components/Icons'
import type { QuickPanelListItem } from '@renderer/components/QuickPanel'
import { QuickPanelReservedSymbol } from '@renderer/components/QuickPanel'
import {
  isGeminiModel,
  isGPT5SeriesReasoningModel,
  isOpenAIWebSearchModel,
  isWebSearchModel
} from '@renderer/config/models'
import { isGeminiWebSearchProvider } from '@renderer/config/providers'
import { useAssistant } from '@renderer/hooks/useAssistant'
import { useTimer } from '@renderer/hooks/useTimer'
import { useWebSearchProviders } from '@renderer/hooks/useWebSearchProviders'
import type { ToolQuickPanelController, ToolRenderContext } from '@renderer/pages/home/Inputbar/types'
import { getProviderByModel } from '@renderer/services/AssistantService'
import WebSearchService from '@renderer/services/WebSearchService'
import type { WebSearchProvider, WebSearchProviderId } from '@renderer/types'
import { hasObjectKey } from '@renderer/utils'
import { isToolUseModeFunction } from '@renderer/utils/assistant'
import { Globe } from 'lucide-react'
import { useCallback, useEffect, useMemo } from 'react'
import { useTranslation } from 'react-i18next'

const logger = loggerService.withContext('WebSearchQuickPanel')

export const WebSearchProviderIcon = ({
  pid,
  size = 18,
  color
}: {
  pid?: WebSearchProviderId
  size?: number
  color?: string
}) => {
  switch (pid) {
    case 'bocha':
      return <BochaLogo className="icon" width={size} height={size} color={color} />
    case 'exa':
      return <ExaLogo className="icon" width={size - 2} height={size} color={color} />
    case 'tavily':
      return <TavilyLogo className="icon" width={size} height={size} color={color} />
    case 'zhipu':
      return <ZhipuLogo className="icon" width={size} height={size} color={color} />
    case 'searxng':
      return <SearXNGLogo className="icon" width={size} height={size} color={color} />
    case 'local-baidu':
      return <BaiduOutlined size={size} style={{ color, fontSize: size }} />
    case 'local-bing':
      return <BingLogo className="icon" width={size} height={size} color={color} />
    case 'local-google':
      return <GoogleOutlined size={size} style={{ color, fontSize: size }} />
    default:
      return <Globe className="icon" size={size} style={{ color, fontSize: size }} />
  }
}

export const useWebSearchPanelController = (assistantId: string, quickPanelController: ToolQuickPanelController) => {
  const { t } = useTranslation()
  const { assistant, updateAssistant } = useAssistant(assistantId)
  const { providers } = useWebSearchProviders()
  const { setTimeoutTimer } = useTimer()

  const enableWebSearch = assistant?.webSearchProviderId || assistant.enableWebSearch

  const updateWebSearchProvider = useCallback(
    async (providerId?: WebSearchProvider['id']) => {
      setTimeoutTimer('updateWebSearchProvider', () => {
        updateAssistant({
          ...assistant,
          webSearchProviderId: providerId,
          enableWebSearch: false
        })
      })
    },
    [assistant, setTimeoutTimer, updateAssistant]
  )

  const updateQuickPanelItem = useCallback(
    async (providerId?: WebSearchProvider['id']) => {
      if (providerId === assistant.webSearchProviderId) {
        updateWebSearchProvider(undefined)
      } else {
        updateWebSearchProvider(providerId)
      }
    },
    [assistant.webSearchProviderId, updateWebSearchProvider]
  )

  const updateToModelBuiltinWebSearch = useCallback(async () => {
    const update = {
      ...assistant,
      webSearchProviderId: undefined,
      enableWebSearch: !assistant.enableWebSearch
    }
    const model = assistant.model
    const provider = getProviderByModel(model)
    if (!model) {
      logger.error('Model does not exist.')
      window.toast.error(t('error.model.not_exists'))
      return
    }
    if (
      isGeminiWebSearchProvider(provider) &&
      isGeminiModel(model) &&
      isToolUseModeFunction(assistant) &&
      update.enableWebSearch &&
      assistant.mcpServers &&
      assistant.mcpServers.length > 0
    ) {
      update.enableWebSearch = false
      window.toast.warning(t('chat.mcp.warning.gemini_web_search'))
    }
    if (
      isOpenAIWebSearchModel(model) &&
      isGPT5SeriesReasoningModel(model) &&
      update.enableWebSearch &&
      assistant.settings?.reasoning_effort === 'minimal'
    ) {
      update.enableWebSearch = false
      window.toast.warning(t('chat.web_search.warning.openai'))
    }
    setTimeoutTimer('updateSelectedWebSearchBuiltin', () => updateAssistant(update), 200)
  }, [assistant, setTimeoutTimer, t, updateAssistant])

  const providerItems = useMemo<QuickPanelListItem[]>(() => {
    const isWebSearchModelEnabled = assistant.model && isWebSearchModel(assistant.model)
    const items: QuickPanelListItem[] = providers
      .map((p) => ({
        label: p.name,
        description: WebSearchService.isWebSearchEnabled(p.id)
          ? hasObjectKey(p, 'apiKey')
            ? t('settings.tool.websearch.apikey')
            : t('settings.tool.websearch.free')
          : t('chat.input.web_search.enable_content'),
        icon: <WebSearchProviderIcon size={13} pid={p.id} />,
        isSelected: p.id === assistant?.webSearchProviderId,
        disabled: !WebSearchService.isWebSearchEnabled(p.id),
        action: () => updateQuickPanelItem(p.id)
      }))
      .filter((item) => !item.disabled)

    if (isWebSearchModelEnabled) {
      items.unshift({
        label: t('chat.input.web_search.builtin.label'),
        description: isWebSearchModelEnabled
          ? t('chat.input.web_search.builtin.enabled_content')
          : t('chat.input.web_search.builtin.disabled_content'),
        icon: <Globe />,
        isSelected: assistant.enableWebSearch,
        disabled: !isWebSearchModelEnabled,
        action: () => updateToModelBuiltinWebSearch()
      })
    }

    return items
  }, [
    assistant.enableWebSearch,
    assistant.model,
    assistant?.webSearchProviderId,
    providers,
    t,
    updateQuickPanelItem,
    updateToModelBuiltinWebSearch
  ])

  const openQuickPanel = useCallback(() => {
    quickPanelController.open({
      title: t('chat.input.web_search.label'),
      list: providerItems,
      symbol: QuickPanelReservedSymbol.WebSearch,
      pageSize: 9
    })
  }, [providerItems, quickPanelController, t])

  const toggleQuickPanel = useCallback(() => {
    if (quickPanelController.isVisible && quickPanelController.symbol === QuickPanelReservedSymbol.WebSearch) {
      quickPanelController.close()
    } else {
      openQuickPanel()
    }
  }, [openQuickPanel, quickPanelController])

  return {
    enableWebSearch,
    providerItems,
    openQuickPanel,
    toggleQuickPanel,
    updateWebSearchProvider,
    updateToModelBuiltinWebSearch,
    selectedProviderId: assistant.webSearchProviderId
  }
}

interface ManagerProps {
  context: ToolRenderContext<any, any>
}

const WebSearchQuickPanelManager = ({ context }: ManagerProps) => {
  const { assistant, quickPanel, quickPanelController, t } = context
  const { providerItems, openQuickPanel } = useWebSearchPanelController(assistant.id, quickPanelController)
  const { registerRootMenu, registerTrigger } = quickPanel
  const { updateList, isVisible, symbol } = quickPanelController

  useEffect(() => {
    if (isVisible && symbol === QuickPanelReservedSymbol.WebSearch) {
      updateList(providerItems)
    }
  }, [isVisible, providerItems, symbol, updateList])

  useEffect(() => {
    const disposeMenu = registerRootMenu([
      {
        label: t('chat.input.web_search.label'),
        description: '',
        icon: <Globe size={18} />,
        isMenu: true,
        action: () => openQuickPanel()
      }
    ])

    const disposeTrigger = registerTrigger(QuickPanelReservedSymbol.WebSearch, () => openQuickPanel())

    return () => {
      disposeMenu()
      disposeTrigger()
    }
  }, [openQuickPanel, registerRootMenu, registerTrigger, t])

  return null
}

export default WebSearchQuickPanelManager