From d388aeecfb518935697f94f297342a03cf6a3454 Mon Sep 17 00:00:00 2001
From: Teo <cheesen.xu@gmail.com>
Date: Tue, 14 Jan 2025 16:38:17 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E6=A8=A1=E5=9E=8B?=
 =?UTF-8?q?=E6=8F=90=E5=8F=8A=E5=8A=9F=E8=83=BD=EF=BC=8C=E6=94=AF=E6=8C=81?=
 =?UTF-8?q?=E5=A4=9A=E4=B8=AA=E6=A8=A1=E5=9E=8B=E4=B8=80=E8=B5=B7=E5=9B=9E?=
 =?UTF-8?q?=E7=AD=94?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../pages/home/Inputbar/AttachmentPreview.tsx |  39 ++---
 .../src/pages/home/Inputbar/Inputbar.tsx      |  36 +++-
 .../home/Inputbar/MentionModelsButton.tsx     | 155 ++++++++++++++++++
 .../home/Inputbar/MentionModelsInput.tsx      |  26 +++
 .../src/pages/home/Messages/Message.tsx       |   4 +-
 .../pages/home/Messages/MessageContent.tsx    |   9 +-
 .../src/pages/home/Messages/Messages.tsx      |  16 +-
 src/renderer/src/types/index.ts               |   2 +
 8 files changed, 254 insertions(+), 33 deletions(-)
 create mode 100644 src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
 create mode 100644 src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx

diff --git a/src/renderer/src/pages/home/Inputbar/AttachmentPreview.tsx b/src/renderer/src/pages/home/Inputbar/AttachmentPreview.tsx
index ba4868d5..a1e7b005 100644
--- a/src/renderer/src/pages/home/Inputbar/AttachmentPreview.tsx
+++ b/src/renderer/src/pages/home/Inputbar/AttachmentPreview.tsx
@@ -16,37 +16,26 @@ const AttachmentPreview: FC<Props> = ({ files, setFiles }) => {
   }
 
   return (
-    <Container>
-      <ContentContainer>
-        <Upload
-          listType={files.length > 20 ? 'text' : 'picture-card'}
-          fileList={files.map((file) => ({
-            uid: file.id,
-            url: 'file://' + FileManager.getSafePath(file),
-            status: 'done',
-            name: file.name
-          }))}
-          onRemove={(item) => setFiles(files.filter((file) => item.uid !== file.id))}
-        />
-      </ContentContainer>
-    </Container>
+    <ContentContainer>
+      <Upload
+        listType={files.length > 20 ? 'text' : 'picture-card'}
+        fileList={files.map((file) => ({
+          uid: file.id,
+          url: 'file://' + FileManager.getSafePath(file),
+          status: 'done',
+          name: file.name
+        }))}
+        onRemove={(item) => setFiles(files.filter((file) => item.uid !== file.id))}
+      />
+    </ContentContainer>
   )
 }
 
-const Container = styled.div`
-  display: flex;
-  flex-direction: row;
-  gap: 10px;
-  padding: 10px 0;
-  background: var(--color-background);
-  border-top: 1px solid var(--color-border-mute);
-`
-
 const ContentContainer = styled.div`
   max-height: 40vh;
-  width: 100%;
   overflow-y: auto;
-  padding: 0 20px;
+  width: 100%;
+  padding: 10px 15px 0;
 `
 
 export default AttachmentPreview
diff --git a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
index 9de9d2a5..0c0e3a9c 100644
--- a/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
+++ b/src/renderer/src/pages/home/Inputbar/Inputbar.tsx
@@ -24,7 +24,7 @@ import { estimateTextTokens as estimateTxtTokens } from '@renderer/services/Toke
 import { translateText } from '@renderer/services/TranslateService'
 import store, { useAppDispatch, useAppSelector } from '@renderer/store'
 import { setGenerating, setSearching } from '@renderer/store/runtime'
-import { Assistant, FileType, KnowledgeBase, Message, Topic } from '@renderer/types'
+import { Assistant, FileType, KnowledgeBase, Message, Model, Topic } from '@renderer/types'
 import { classNames, delay, getFileExtension, uuid } from '@renderer/utils'
 import { documentExts, imageExts, textExts } from '@shared/config/constant'
 import { Button, Popconfirm, Tooltip } from 'antd'
@@ -39,6 +39,8 @@ import NarrowLayout from '../Messages/NarrowLayout'
 import AttachmentButton from './AttachmentButton'
 import AttachmentPreview from './AttachmentPreview'
 import KnowledgeBaseButton from './KnowledgeBaseButton'
+import MentionModelsButton from './MentionModelsButton'
+import MentionModelsInput from './MentionModelsInput'
 import SendMessageButton from './SendMessageButton'
 import TokenCount from './TokenCount'
 
@@ -82,6 +84,7 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
   const spaceClickTimer = useRef<NodeJS.Timeout>()
   const [isTranslating, setIsTranslating] = useState(false)
   const [selectedKnowledgeBase, setSelectedKnowledgeBase] = useState<KnowledgeBase | undefined>(_base)
+  const [mentionModels, setMentionModels] = useState<Model[]>([])
 
   const isVision = useMemo(() => isVisionModel(model), [model])
   const supportExts = useMemo(() => [...textExts, ...documentExts, ...(isVision ? imageExts : [])], [isVision])
@@ -126,15 +129,20 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
       message.files = await FileManager.uploadFiles(files)
     }
 
+    if (mentionModels.length > 0) {
+      message.mentions = mentionModels
+    }
+
     EventEmitter.emit(EVENT_NAMES.SEND_MESSAGE, message)
 
     setText('')
     setFiles([])
+    setMentionModels([])
     setTimeout(() => setText(''), 500)
     setTimeout(() => resizeTextArea(), 0)
 
     setExpend(false)
-  }, [inputEmpty, text, assistant.id, assistant.topics, selectedKnowledgeBase, files])
+  }, [inputEmpty, text, assistant.id, assistant.topics, selectedKnowledgeBase, files, mentionModels])
 
   const translate = async () => {
     if (isTranslating) {
@@ -386,14 +394,31 @@ 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])
+      }
+    },
+    [mentionModels]
+  )
+
+  const handleRemoveModel = (model: Model) => {
+    setMentionModels(mentionModels.filter((m) => m.id !== model.id))
+  }
+
   return (
     <Container onDragOver={handleDragOver} onDrop={handleDrop} className="inputbar">
       <NarrowLayout style={{ width: '100%' }}>
-        <AttachmentPreview files={files} setFiles={setFiles} />
         <InputBarContainer
           id="inputbar"
           className={classNames('inputbar-container', inputFocus && 'focus')}
           ref={containerRef}>
+          <AttachmentPreview files={files} setFiles={setFiles} />
+          <MentionModelsInput selectedModels={mentionModels} onRemoveModel={handleRemoveModel} />
           <Textarea
             value={text}
             onChange={(e) => setText(e.target.value)}
@@ -421,6 +446,11 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic }) => {
                   <FormOutlined />
                 </ToolbarButton>
               </Tooltip>
+              <MentionModelsButton
+                mentionModels={mentionModels}
+                onMentionModel={onMentionModel}
+                ToolbarButton={ToolbarButton}
+              />
               {isWebSearchModel(model) && (
                 <Tooltip placement="top" title={t('chat.input.web_search')} arrow>
                   <ToolbarButton
diff --git a/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx b/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
new file mode 100644
index 00000000..11fe5e4d
--- /dev/null
+++ b/src/renderer/src/pages/home/Inputbar/MentionModelsButton.tsx
@@ -0,0 +1,155 @@
+import { PushpinOutlined } from '@ant-design/icons'
+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 { getModelUniqId } from '@renderer/services/ModelService'
+import { Model } from '@renderer/types'
+import { Avatar, Dropdown, Tooltip } from 'antd'
+import { first, sortBy } from 'lodash'
+import { FC, useEffect, useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import styled, { createGlobalStyle } from 'styled-components'
+
+interface Props {
+  mentionModels: Model[]
+  onMentionModel: (model: Model) => void
+  ToolbarButton: any
+}
+
+const MentionModelsButton: FC<Props> = ({ 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 togglePin = async (modelId: string) => {
+    const newPinnedModels = pinnedModels.includes(modelId)
+      ? pinnedModels.filter((id) => id !== modelId)
+      : [...pinnedModels, modelId]
+
+    await db.settings.put({ id: 'pinned:models', value: newPinnedModels })
+    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))
+        .map((m) => ({
+          key: getModelUniqId(m),
+          label: (
+            <ModelItem>
+              <span>
+                {m?.name} <ModelTags model={m} />
+              </span>
+              {/* <Checkbox checked={selectedModels.some((sm) => sm.id === m.id)} /> */}
+              <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: () => {
+            onSelect(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(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
+      })
+    }
+  }
+
+  return (
+    <>
+      <DropdownMenuStyle />
+      <Dropdown menu={{ items: modelMenuItems }} trigger={['click']} overlayClassName="mention-models-dropdown">
+        <Tooltip placement="top" title={t('agents.edit.model.select.title')} arrow>
+          <ToolbarButton type="text">
+            <i className="iconfont icon-at"></i>
+          </ToolbarButton>
+        </Tooltip>
+      </Dropdown>
+    </>
+  )
+}
+
+const DropdownMenuStyle = createGlobalStyle`
+  .mention-models-dropdown {
+    .ant-dropdown-menu {
+      max-height: 400px;
+    }
+  }
+`
+
+const ModelItem = styled.div`
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  font-size: 14px;
+  width: 100%;
+  gap: 16px;
+
+  &:hover {
+    .pin-icon {
+      opacity: 0.3;
+    }
+  }
+`
+
+const PinIcon = styled.span.attrs({ className: 'pin-icon' })<{ $isPinned: boolean }>`
+  margin-left: auto;
+  padding: 0 8px;
+  opacity: ${(props) => (props.$isPinned ? 1 : 'inherit')};
+  transition: opacity 0.2s;
+  right: 0;
+  color: ${(props) => (props.$isPinned ? 'var(--color-primary)' : 'inherit')};
+  transform: ${(props) => (props.$isPinned ? 'rotate(-45deg)' : 'none')};
+  opacity: 0;
+
+  &:hover {
+    opacity: 1 !important;
+    color: ${(props) => (props.$isPinned ? 'var(--color-primary)' : 'inherit')};
+  }
+`
+
+export default MentionModelsButton
diff --git a/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx b/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx
new file mode 100644
index 00000000..f2e44fd2
--- /dev/null
+++ b/src/renderer/src/pages/home/Inputbar/MentionModelsInput.tsx
@@ -0,0 +1,26 @@
+import { Model } from '@renderer/types'
+import { Flex, Tag } from 'antd'
+import { FC } from 'react'
+import styled from 'styled-components'
+
+const MentionModelsInput: FC<{
+  selectedModels: Model[]
+  onRemoveModel: (model: Model) => void
+}> = ({ selectedModels, onRemoveModel }) => {
+  return (
+    <Container gap="4px 0" wrap>
+      {selectedModels.map((model) => (
+        <Tag bordered={false} color="processing" key={model.id} closable onClose={() => onRemoveModel(model)}>
+          @{model.name}
+        </Tag>
+      ))}
+    </Container>
+  )
+}
+
+const Container = styled(Flex)`
+  width: 100%;
+  padding: 10px 15px 0;
+`
+
+export default MentionModelsInput
diff --git a/src/renderer/src/pages/home/Messages/Message.tsx b/src/renderer/src/pages/home/Messages/Message.tsx
index 9aa80940..53ba8a23 100644
--- a/src/renderer/src/pages/home/Messages/Message.tsx
+++ b/src/renderer/src/pages/home/Messages/Message.tsx
@@ -109,6 +109,8 @@ const MessageItem: FC<Props> = ({
     if (topic && onGetMessages && onSetMessages) {
       if (message.status === 'sending') {
         const messages = onGetMessages()
+        const assistantWithModel = message.model ? { ...assistant, model: message.model } : assistant
+
         fetchChatCompletion({
           message,
           messages: messages
@@ -117,7 +119,7 @@ const MessageItem: FC<Props> = ({
               0,
               messages.findIndex((m) => m.id === message.id)
             ),
-          assistant,
+          assistant: assistantWithModel,
           topic,
           onResponse: (msg) => {
             setMessage(msg)
diff --git a/src/renderer/src/pages/home/Messages/MessageContent.tsx b/src/renderer/src/pages/home/Messages/MessageContent.tsx
index 5bde1192..769e6a41 100644
--- a/src/renderer/src/pages/home/Messages/MessageContent.tsx
+++ b/src/renderer/src/pages/home/Messages/MessageContent.tsx
@@ -1,7 +1,7 @@
 import { SyncOutlined, TranslationOutlined } from '@ant-design/icons'
 import { Message, Model } from '@renderer/types'
 import { getBriefInfo } from '@renderer/utils'
-import { Divider } from 'antd'
+import { Divider, Flex } from 'antd'
 import React from 'react'
 import { useTranslation } from 'react-i18next'
 import BeatLoader from 'react-spinners/BeatLoader'
@@ -37,6 +37,9 @@ const MessageContent: React.FC<{
 
   return (
     <>
+      <Flex gap="8px" wrap>
+        {message.mentions?.map((model) => <MentionTag key={model.id}>{'@' + model.name}</MentionTag>)}
+      </Flex>
       <Markdown message={message} />
       {message.translatedContent && (
         <>
@@ -65,4 +68,8 @@ const MessageContentLoading = styled.div`
   margin-bottom: 5px;
 `
 
+const MentionTag = styled.span`
+  color: var(--color-link);
+`
+
 export default React.memo(MessageContent)
diff --git a/src/renderer/src/pages/home/Messages/Messages.tsx b/src/renderer/src/pages/home/Messages/Messages.tsx
index 0f9023ec..6d32b358 100644
--- a/src/renderer/src/pages/home/Messages/Messages.tsx
+++ b/src/renderer/src/pages/home/Messages/Messages.tsx
@@ -97,10 +97,19 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
 
   const onSendMessage = useCallback(
     async (message: Message) => {
-      const assistantMessage = getAssistantMessage({ assistant, topic })
+      const assistantMessages: Message[] = []
+      if (message.mentions?.length) {
+        message.mentions.forEach((m) => {
+          const assistantMessage = getAssistantMessage({ assistant: { ...assistant, model: m }, topic })
+          assistantMessage.model = m
+          assistantMessages.push(assistantMessage)
+        })
+      } else {
+        assistantMessages.push(getAssistantMessage({ assistant, topic }))
+      }
 
       setMessages((prev) => {
-        const messages = prev.concat([message, assistantMessage])
+        const messages = prev.concat([message, ...assistantMessages])
         db.topics.put({ id: topic.id, messages })
         return messages
       })
@@ -156,7 +165,8 @@ const Messages: FC<Props> = ({ assistant, topic, setActiveTopic }) => {
       }),
       EventEmitter.on(EVENT_NAMES.REGENERATE_MESSAGE, async (model: Model) => {
         const lastUserMessage = last(filterMessages(messages).filter((m) => m.role === 'user'))
-        lastUserMessage && onSendMessage({ ...lastUserMessage, id: uuid(), type: '@', modelId: model.id })
+        lastUserMessage &&
+          onSendMessage({ ...lastUserMessage, id: uuid(), modelId: model.id, model: model, mentions: [model] })
       }),
       EventEmitter.on(EVENT_NAMES.AI_AUTO_RENAME, autoRenameTopic),
       EventEmitter.on(EVENT_NAMES.CLEAR_MESSAGES, () => {
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index 1297ee2a..78d7da88 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -59,6 +59,8 @@ export type Message = {
   knowledgeBaseIds?: string[]
   type: 'text' | '@' | 'clear'
   isPreset?: boolean
+  mentions?: Model[]
+  model?: Model
   metadata?: {
     // Gemini
     groundingMetadata?: any