From 5e8d7682f5f75893e7f88c59fd623b2dd1e56315 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=A2=E5=A5=8B=E7=8C=AB?= <kangfenmao@qq.com>
Date: Wed, 12 Feb 2025 14:08:18 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20=E4=BD=BF=E7=94=A8@?=
 =?UTF-8?q?=E5=91=BC=E5=87=BA=E6=A8=A1=E5=9E=8B=E9=80=89=E6=8B=A9=E5=88=97?=
 =?UTF-8?q?=E8=A1=A8=20#1317=20#1324=20(#1458)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* feat: 🎸 使用@呼出模型选择列表

输入第一个字符为@符号的时候可以呼出选择模型的列表

* feat: 🎸 Only one can be chosen at a time

一次只能选择一个模型，选择后自动关闭。选择过的模型不在出现在列表，避免删除模型的时候显示异常。

* fix: 🐛 When choosing the model, Enter will send a message

* feat: 🎸 选中的模型显示供应商

* feat: 🎸 pinned module show privoder

* feat: 🎸 only selected modle show provider

* feat: 🎸 删除@符号以后自动关闭

* feat: 🎸 增加模糊搜索

---------

Co-authored-by: duanyongcheng77 <duanyongcheng77@gmail.com>
---
 .../src/pages/home/Inputbar/Inputbar.tsx      |  78 ++++-
 .../home/Inputbar/MentionModelsButton.tsx     | 325 +++++++++++++++---
 .../home/Inputbar/MentionModelsInput.tsx      |  12 +-
 src/renderer/src/services/EventService.ts     |   3 +-
 4 files changed, 350 insertions(+), 68 deletions(-)

diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
index 4f89b148..9b441596 100644
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@@ -85,6 +85,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
   const [isTranslating, setIsTranslating] = useState(false)
   const [selectedKnowledgeBase, setSelectedKnowledgeBase] = useState<KnowledgeBase | undefined>(_base)
   const [mentionModels, setMentionModels] = useState<Model[]>([])
+  const [isMentionPopupOpen, setIsMentionPopupOpen] = useState(false)
 
   const isVision = useMemo(() => isVisionModel(model), [model])
   const supportExts = useMemo(() => [...textExts, ...documentExts, ...(isVision ? imageExts : [])], [isVision])
@@ -165,6 +166,24 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
   const handleKeyDown = (event: React.KeyboardEvent<HTMLTextAreaElement>) => {
     const isEnterPressed = event.keyCode == 13
 
+    if (event.key === '@') {
+      const textArea = textareaRef.current?.resizableTextArea?.textArea
+      if (textArea) {
+        const cursorPosition = textArea.selectionStart
+        const textBeforeCursor = text.substring(0, cursorPosition)
+        if (cursorPosition === 0 || textBeforeCursor.endsWith(' ')) {
+          EventEmitter.emit(EVENT_NAMES.SHOW_MODEL_SELECTOR)
+          setIsMentionPopupOpen(true)
+          return
+        }
+      }
+    }
+
+    if (event.key === 'Escape' && isMentionPopupOpen) {
+      setIsMentionPopupOpen(false)
+      return
+    }
+
     if (autoTranslateWithSpace) {
       if (event.key === ' ') {
         setSpaceClickCount((prev) => prev + 1)
@@ -193,25 +212,34 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
       }
     }
 
-    if (sendMessageShortcut === 'Enter' && isEnterPressed) {
-      if (event.shiftKey) {
-        return
+    if (isEnterPressed && !event.shiftKey && sendMessageShortcut === 'Enter') {
+      if (isMentionPopupOpen) {
+        return event.preventDefault()
       }
       sendMessage()
       return event.preventDefault()
     }
 
     if (sendMessageShortcut === 'Shift+Enter' && isEnterPressed && event.shiftKey) {
+      if (isMentionPopupOpen) {
+        return event.preventDefault()
+      }
       sendMessage()
       return event.preventDefault()
     }
 
     if (sendMessageShortcut === 'Ctrl+Enter' && isEnterPressed && event.ctrlKey) {
+      if (isMentionPopupOpen) {
+        return event.preventDefault()
+      }
       sendMessage()
       return event.preventDefault()
     }
 
     if (sendMessageShortcut === 'Command+Enter' && isEnterPressed && event.metaKey) {
+      if (isMentionPopupOpen) {
+        return event.preventDefault()
+      }
       sendMessage()
       return event.preventDefault()
     }
@@ -280,6 +308,23 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
 
   const onInput = () => !expended && resizeTextArea()
 
+  const onChange = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
+    const newText = e.target.value
+    setText(newText)
+
+    // Check if @ was deleted
+    const textArea = textareaRef.current?.resizableTextArea?.textArea
+    if (textArea) {
+      const cursorPosition = textArea.selectionStart
+      const textBeforeCursor = newText.substring(0, cursorPosition)
+      const lastAtIndex = textBeforeCursor.lastIndexOf('@')
+
+      if (lastAtIndex === -1 || textBeforeCursor.slice(lastAtIndex + 1).includes(' ')) {
+        setIsMentionPopupOpen(false)
+      }
+    }
+  }
+
   const onPaste = useCallback(
     async (event: ClipboardEvent) => {
       const clipboardText = event.clipboardData?.getData('text')
@@ -420,17 +465,22 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
     setSelectedKnowledgeBase(base)
   }
 
-  const onMentionModel = useCallback(
-    (model: Model) => {
-      const isSelected = mentionModels.some((m) => m.id === model.id)
-      if (isSelected) {
-        setMentionModels(mentionModels.filter((m) => m.id !== model.id))
-      } else {
-        setMentionModels([...mentionModels, model])
+  const onMentionModel = (model: Model) => {
+    const textArea = textareaRef.current?.resizableTextArea?.textArea
+    if (textArea) {
+      const cursorPosition = textArea.selectionStart
+      const textBeforeCursor = text.substring(0, cursorPosition)
+      const lastAtIndex = textBeforeCursor.lastIndexOf('@')
+
+      if (lastAtIndex !== -1) {
+        const newText = text.substring(0, lastAtIndex) + text.substring(cursorPosition)
+        setText(newText)
       }
-    },
-    [mentionModels]
-  )
+
+      setMentionModels((prev) => [...prev, model])
+      setIsMentionPopupOpen(false)
+    }
+  }
 
   const handleRemoveModel = (model: Model) => {
     setMentionModels(mentionModels.filter((m) => m.id !== model.id))
@@ -447,7 +497,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
           <MentionModelsInput selectedModels={mentionModels} onRemoveModel={handleRemoveModel} />
           <Textarea
             value={text}
-            onChange={(e) => setText(e.target.value)}
+            onChange={onChange}
             onKeyDown={handleKeyDown}
             placeholder={isTranslating ? t('chat.input.translating') : t('chat.input.placeholder')}
             autoFocus
diff --git a/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx b/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
index 05a9c796..92f1f5e1 100644
--- a/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
+++ b/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
@@ -3,11 +3,12 @@ import ModelTags from '@renderer/components/ModelTags'
 import { getModelLogo, isEmbeddingModel } from '@renderer/config/models'
 import db from '@renderer/databases'
 import { useProviders } from '@renderer/hooks/useProvider'
+import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
 import { getModelUniqId } from '@renderer/services/ModelService'
-import { Model } from '@renderer/types'
+import { Model, Provider } from '@renderer/types'
 import { Avatar, Dropdown, Tooltip } from 'antd'
 import { first, sortBy } from 'lodash'
-import { FC, useEffect, useState } from 'react'
+import { FC, useEffect, useMemo, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled, { createGlobalStyle } from 'styled-components'
 
@@ -17,18 +18,15 @@ interface Props {
   ToolbarButton: any
 }
 
-const MentionModelsButton: FC<Props> = ({ onMentionModel: onSelect, ToolbarButton }) => {
+const MentionModelsButton: FC<Props> = ({ mentionModels, onMentionModel: onSelect, ToolbarButton }) => {
   const { providers } = useProviders()
   const [pinnedModels, setPinnedModels] = useState<string[]>([])
   const { t } = useTranslation()
-
-  useEffect(() => {
-    const loadPinnedModels = async () => {
-      const setting = await db.settings.get('pinned:models')
-      setPinnedModels(setting?.value || [])
-    }
-    loadPinnedModels()
-  }, [])
+  const dropdownRef = useRef<any>(null)
+  const [selectedIndex, setSelectedIndex] = useState(-1)
+  const [isOpen, setIsOpen] = useState(false)
+  const menuRef = useRef<HTMLDivElement>(null)
+  const [searchText, setSearchText] = useState('')
 
   const togglePin = async (modelId: string) => {
     const newPinnedModels = pinnedModels.includes(modelId)
@@ -39,72 +37,246 @@ const MentionModelsButton: FC<Props> = ({ onMentionModel: onSelect, ToolbarButto
     setPinnedModels(newPinnedModels)
   }
 
-  const modelMenuItems = providers
-    .filter((p) => p.models && p.models.length > 0)
-    .map((p) => {
-      const filteredModels = sortBy(p.models, ['group', 'name'])
-        .filter((m) => !isEmbeddingModel(m))
+  const handleModelSelect = (model: Model) => {
+    // Check if model is already selected
+    if (mentionModels.some((selected) => selected.id === model.id)) {
+      return
+    }
+    onSelect(model)
+    setIsOpen(false)
+  }
+
+  const modelMenuItems = useMemo(() => {
+    const items = providers
+      .filter((p) => p.models && p.models.length > 0)
+      .map((p) => {
+        const filteredModels = sortBy(p.models, ['group', 'name'])
+          .filter((m) => !isEmbeddingModel(m))
+          // Filter out pinned models from regular groups
+          .filter((m) => !pinnedModels.includes(getModelUniqId(m)))
+          // Filter by search text
+          .filter((m) => {
+            if (!searchText) return true
+            return (
+              m.name.toLowerCase().includes(searchText.toLowerCase()) ||
+              m.id.toLowerCase().includes(searchText.toLowerCase())
+            )
+          })
+          .map((m) => ({
+            key: getModelUniqId(m),
+            model: m,
+            label: (
+              <ModelItem>
+                <ModelNameRow>
+                  <span>{m?.name}</span> <ModelTags model={m} />
+                </ModelNameRow>
+                <PinIcon
+                  onClick={(e) => {
+                    e.stopPropagation()
+                    togglePin(getModelUniqId(m))
+                  }}
+                  $isPinned={pinnedModels.includes(getModelUniqId(m))}>
+                  <PushpinOutlined />
+                </PinIcon>
+              </ModelItem>
+            ),
+            icon: (
+              <Avatar src={getModelLogo(m.id)} size={24}>
+                {first(m.name)}
+              </Avatar>
+            ),
+            onClick: () => handleModelSelect(m)
+          }))
+
+        return filteredModels.length > 0
+          ? {
+              key: p.id,
+              label: p.isSystem ? t(`provider.${p.id}`) : p.name,
+              type: 'group' as const,
+              children: filteredModels
+            }
+          : null
+      })
+      .filter((group): group is NonNullable<typeof group> => group !== null)
+
+    if (pinnedModels.length > 0) {
+      const pinnedItems = providers
+        .filter((p): p is Provider => p.models && p.models.length > 0)
+        .flatMap((p) =>
+          p.models
+            .filter((m) => pinnedModels.includes(getModelUniqId(m)))
+            .map((m) => ({
+              key: getModelUniqId(m),
+              model: m,
+              provider: p
+            }))
+        )
         .map((m) => ({
-          key: getModelUniqId(m),
+          ...m,
+          key: m.key + 'pinned',
           label: (
             <ModelItem>
               <ModelNameRow>
-                <span>{m?.name}</span> <ModelTags model={m} />
+                <span>
+                  {m.model?.name} | {m.provider.isSystem ? t(`provider.${m.provider.id}`) : m.provider.name}
+                </span>{' '}
+                <ModelTags model={m.model} />
               </ModelNameRow>
-              {/* <Checkbox checked={selectedModels.some((sm) => sm.id === m.id)} /> */}
               <PinIcon
                 onClick={(e) => {
                   e.stopPropagation()
-                  togglePin(getModelUniqId(m))
+                  togglePin(getModelUniqId(m.model))
                 }}
-                $isPinned={pinnedModels.includes(getModelUniqId(m))}>
+                $isPinned={true}>
                 <PushpinOutlined />
               </PinIcon>
             </ModelItem>
           ),
           icon: (
-            <Avatar src={getModelLogo(m.id)} size={24}>
-              {first(m.name)}
+            <Avatar src={getModelLogo(m.model.id)} size={24}>
+              {first(m.model.name)}
             </Avatar>
           ),
-          onClick: () => {
-            onSelect(m)
-          }
+          onClick: () => handleModelSelect(m.model)
         }))
 
-      return filteredModels.length > 0
-        ? {
-            key: p.id,
-            label: p.isSystem ? t(`provider.${p.id}`) : p.name,
-            type: 'group' as const,
-            children: filteredModels
-          }
-        : null
-    })
-    .filter(Boolean)
-
-  if (pinnedModels.length > 0) {
-    const pinnedItems = modelMenuItems
-      .flatMap((p) => p?.children || [])
-      .filter((m) => pinnedModels.includes(m.key))
-      .map((m) => ({ ...m, key: m.key + 'pinned' }))
-
-    if (pinnedItems.length > 0) {
-      modelMenuItems.unshift({
-        key: 'pinned',
-        label: t('models.pinned'),
-        type: 'group' as const,
-        children: pinnedItems
-      })
+      if (pinnedItems.length > 0) {
+        items.unshift({
+          key: 'pinned',
+          label: t('models.pinned'),
+          type: 'group' as const,
+          children: pinnedItems
+        })
+      }
     }
+
+    // Remove empty groups
+    return items.filter((group) => group.children.length > 0)
+  }, [providers, pinnedModels, t, onSelect, mentionModels, searchText])
+
+  // Get flattened list of all model items
+  const flatModelItems = useMemo(() => {
+    return modelMenuItems.flatMap((group) => group?.children || [])
+  }, [modelMenuItems])
+
+  useEffect(() => {
+    const loadPinnedModels = async () => {
+      const setting = await db.settings.get('pinned:models')
+      setPinnedModels(setting?.value || [])
+    }
+    loadPinnedModels()
+  }, [])
+
+  useEffect(() => {
+    const showModelSelector = () => {
+      dropdownRef.current?.click()
+      setIsOpen(true)
+      setSelectedIndex(0)
+      setSearchText('')
+    }
+
+    const handleKeyDown = (e: KeyboardEvent) => {
+      if (!isOpen) return
+
+      if (e.key === 'ArrowDown') {
+        e.preventDefault()
+        setSelectedIndex((prev) => (prev < flatModelItems.length - 1 ? prev + 1 : prev))
+      } else if (e.key === 'ArrowUp') {
+        e.preventDefault()
+        setSelectedIndex((prev) => (prev > 0 ? prev - 1 : prev))
+      } else if (e.key === 'Enter') {
+        e.preventDefault()
+        if (selectedIndex >= 0 && selectedIndex < flatModelItems.length) {
+          const selectedModel = flatModelItems[selectedIndex].model
+          if (!mentionModels.some((selected) => selected.id === selectedModel.id)) {
+            flatModelItems[selectedIndex].onClick()
+          }
+          setIsOpen(false)
+          setSearchText('')
+        }
+      } else if (e.key === 'Escape') {
+        setIsOpen(false)
+        setSearchText('')
+      }
+    }
+
+    const handleTextChange = (e: Event) => {
+      const textArea = e.target as HTMLTextAreaElement
+      const cursorPosition = textArea.selectionStart
+      const textBeforeCursor = textArea.value.substring(0, cursorPosition)
+      const lastAtIndex = textBeforeCursor.lastIndexOf('@')
+
+      if (lastAtIndex === -1 || textBeforeCursor.slice(lastAtIndex + 1).includes(' ')) {
+        setIsOpen(false)
+        setSearchText('')
+      } else if (lastAtIndex !== -1) {
+        // Get the text after @ for search
+        const searchStr = textBeforeCursor.slice(lastAtIndex + 1)
+        setSearchText(searchStr)
+      }
+    }
+
+    const textArea = document.querySelector('.inputbar textarea') as HTMLTextAreaElement
+    if (textArea) {
+      textArea.addEventListener('input', handleTextChange)
+    }
+
+    EventEmitter.on(EVENT_NAMES.SHOW_MODEL_SELECTOR, showModelSelector)
+    document.addEventListener('keydown', handleKeyDown)
+
+    return () => {
+      EventEmitter.off(EVENT_NAMES.SHOW_MODEL_SELECTOR, showModelSelector)
+      document.removeEventListener('keydown', handleKeyDown)
+      if (textArea) {
+        textArea.removeEventListener('input', handleTextChange)
+      }
+    }
+  }, [isOpen, selectedIndex, flatModelItems, mentionModels])
+
+  // Hide dropdown if no models available
+  if (flatModelItems.length === 0) {
+    return null
   }
 
+  const menu = (
+    <div ref={menuRef} className="ant-dropdown-menu">
+      {modelMenuItems.map((group, groupIndex) => {
+        if (!group) return null
+
+        // Calculate the starting index for this group's items
+        const startIndex = modelMenuItems.slice(0, groupIndex).reduce((acc, g) => acc + (g?.children?.length || 0), 0)
+
+        return (
+          <div key={group.key} className="ant-dropdown-menu-item-group">
+            <div className="ant-dropdown-menu-item-group-title">{group.label}</div>
+            <div>
+              {group.children.map((item, idx) => (
+                <div
+                  key={item.key}
+                  className={`ant-dropdown-menu-item ${selectedIndex === startIndex + idx ? 'ant-dropdown-menu-item-selected' : ''}`}
+                  onClick={item.onClick}>
+                  <span className="ant-dropdown-menu-item-icon">{item.icon}</span>
+                  {item.label}
+                </div>
+              ))}
+            </div>
+          </div>
+        )
+      })}
+    </div>
+  )
+
   return (
     <>
       <DropdownMenuStyle />
-      <Dropdown menu={{ items: modelMenuItems }} trigger={['click']} overlayClassName="mention-models-dropdown">
+      <Dropdown
+        dropdownRender={() => menu}
+        trigger={['click']}
+        open={isOpen}
+        onOpenChange={setIsOpen}
+        overlayClassName="mention-models-dropdown">
         <Tooltip placement="top" title={t('agents.edit.model.select.title')} arrow>
-          <ToolbarButton type="text">
+          <ToolbarButton type="text" ref={dropdownRef}>
             <i className="iconfont icon-at" style={{ fontSize: 18 }}></i>
           </ToolbarButton>
         </Tooltip>
@@ -117,6 +289,54 @@ const DropdownMenuStyle = createGlobalStyle`
   .mention-models-dropdown {
     .ant-dropdown-menu {
       max-height: 400px;
+      overflow-y: auto;
+      overflow-x: hidden;
+      padding: 4px 0;
+      margin-bottom: 40px;
+
+      &::-webkit-scrollbar {
+        width: 6px;
+        height: 6px;
+      }
+
+      &::-webkit-scrollbar-thumb {
+        background: var(--color-scrollbar);
+        border-radius: 3px;
+      }
+
+      &::-webkit-scrollbar-track {
+        background: transparent;
+      }
+    }
+
+    .ant-dropdown-menu-item-group {
+      .ant-dropdown-menu-item-group-title {
+        padding: 5px 12px;
+        color: var(--color-text-3);
+        font-size: 12px;
+      }
+    }
+
+    .ant-dropdown-menu-item {
+      padding: 5px 12px;
+      cursor: pointer;
+      transition: all 0.3s;
+      display: flex;
+      align-items: center;
+      gap: 8px;
+
+      &:hover {
+        background: var(--color-hover);
+      }
+
+      &.ant-dropdown-menu-item-selected {
+        background-color: var(--color-primary-bg);
+        color: var(--color-primary);
+      }
+
+      .ant-dropdown-menu-item-icon {
+        margin-right: 0;
+      }
     }
   }
 `
@@ -127,6 +347,7 @@ const ModelItem = styled.div`
   justify-content: space-between;
   font-size: 14px;
   width: 100%;
+  min-width: 200px;
   gap: 16px;
 
   &:hover {
diff --git a/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx b/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx
index f2e44fd2..9c91d97c 100644
--- a/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx
+++ b/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx
@@ -1,17 +1,27 @@
+import { useProviders } from '@renderer/hooks/useProvider'
 import { Model } from '@renderer/types'
 import { Flex, Tag } from 'antd'
 import { FC } from 'react'
+import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
 
 const MentionModelsInput: FC<{
   selectedModels: Model[]
   onRemoveModel: (model: Model) => void
 }> = ({ selectedModels, onRemoveModel }) => {
+  const { providers } = useProviders()
+  const { t } = useTranslation()
+
+  const getProviderName = (model: Model) => {
+    const provider = providers.find((p) => p.models?.some((m) => m.id === model.id))
+    return provider ? (provider.isSystem ? t(`provider.${provider.id}`) : provider.name) : ''
+  }
+
   return (
     <Container gap="4px 0" wrap>
       {selectedModels.map((model) => (
         <Tag bordered={false} color="processing" key={model.id} closable onClose={() => onRemoveModel(model)}>
-          @{model.name}
+          @{model.name} ({getProviderName(model)})
         </Tag>
       ))}
     </Container>
diff --git a/src/renderer/src/services/EventService.ts b/src/renderer/src/services/EventService.ts
index f31c946a..d8901b83 100644
--- a/src/renderer/src/services/EventService.ts
+++ b/src/renderer/src/services/EventService.ts
@@ -22,5 +22,6 @@ export const EVENT_NAMES = {
   EXPORT_TOPIC_IMAGE: 'EXPORT_TOPIC_IMAGE',
   LOCATE_MESSAGE: 'LOCATE_MESSAGE',
   ADD_NEW_TOPIC: 'ADD_NEW_TOPIC',
-  RESEND_MESSAGE: 'RESEND_MESSAGE'
+  RESEND_MESSAGE: 'RESEND_MESSAGE',
+  SHOW_MODEL_SELECTOR: 'SHOW_MODEL_SELECTOR'
 }