feat: add grounding source info to gemini message

src/renderer/src/assets/styles/index.scss

@@ -42,6 +42,9 @@
   --color-active: rgba(55, 55, 55, 1);
   --color-frame-border: #333;
   --color-group-background: var(--color-background-soft);
+
+  --color-reference: #404040;
+  --color-reference-text: #ffffff;
   --color-reference-background: #0b0e12;
 
   --navbar-background-mac: rgba(30, 30, 30, 0.6);
@@ -102,6 +105,9 @@ body[theme-mode='light'] {
   --color-active: var(--color-white-soft);
   --color-frame-border: #ddd;
   --color-group-background: var(--color-white);
+
+  --color-reference: #cfe1ff;
+  --color-reference-text: #000000;
   --color-reference-background: #f1f7ff;
 
   --navbar-background-mac: rgba(255, 255, 255, 0.6);

src/renderer/src/assets/styles/markdown.scss

@@ -208,6 +208,14 @@
 
   sup {
     top: -0.5em;
+    border-radius: 50%;
+    background-color: var(--color-reference);
+    color: var(--color-reference-text);
+    padding: 2px 5px;
+    zoom: 0.8;
+    & > span.link {
+      color: var(--color-white);
+    }
   }
 
   sub {
@@ -226,6 +234,7 @@
       text-decoration: underline;
     }
   }
+}
 
 .footnotes {
   margin-top: 1em;
@@ -241,6 +250,10 @@
     font-size: 12px;
   }
 
+  a {
+    color: var(--color-link);
+  }
+
   ol {
     padding-left: 1em;
     margin: 0;
@@ -273,7 +286,6 @@
     }
   }
 }
-}
 
 emoji-picker {
   --border-size: 0;

src/renderer/src/config/models.ts

@@ -123,7 +123,6 @@ import YiModelLogo from '@renderer/assets/images/models/yi.png'
 import YiModelLogoDark from '@renderer/assets/images/models/yi_dark.png'
 import { getProviderByModel } from '@renderer/services/AssistantService'
 import { Model } from '@renderer/types'
-import { isEmpty } from 'lodash'
 import OpenAI from 'openai'
 
 import { getWebSearchTools } from './tools'
@@ -1087,28 +1086,14 @@ export function isWebSearchModel(model: Model): boolean {
 
 export function getOpenAIWebSearchParams(model: Model): Record<string, any> {
   if (isWebSearchModel(model)) {
+    const webSearchTools = getWebSearchTools(model)
+
     if (model.provider === 'hunyuan') {
       return { enable_enhancement: true }
     }
 
-    if (model.provider === 'zhipu') {
-      const webSearchTools = getWebSearchTools(model)
-      return isEmpty(webSearchTools)
-        ? {}
-        : {
-            tools: webSearchTools
-          }
-    }
-
     return {
-      tools: [
-        {
-          type: 'function',
-          function: {
-            name: 'googleSearch'
-          }
-        }
-      ]
+      tools: webSearchTools
     }
   }
 

src/renderer/src/config/tools.ts

@@ -1,11 +1,8 @@
 import { Model } from '@renderer/types'
 import { ChatCompletionTool } from 'openai/resources'
 
-import { isWebSearchModel } from './models'
-
 export function getWebSearchTools(model: Model): ChatCompletionTool[] {
-  if (model && model.provider === 'zhipu') {
-    if (isWebSearchModel(model)) {
+  if (model?.provider === 'zhipu') {
     if (model.id === 'glm-4-alltools') {
       return [
         {
@@ -23,7 +20,13 @@ export function getWebSearchTools(model: Model): ChatCompletionTool[] {
       } as unknown as ChatCompletionTool
     ]
   }
-  }
 
-  return []
+  return [
+    {
+      type: 'function',
+      function: {
+        name: 'googleSearch'
+      }
+    }
+  ]
 }

src/renderer/src/pages/home/Markdown/Artifacts.tsx

@@ -1,7 +1,7 @@
 import { DownloadOutlined, ExpandOutlined } from '@ant-design/icons'
 import MinApp from '@renderer/components/MinApp'
 import { AppLogo } from '@renderer/config/env'
-import { extractTitle } from '@renderer/utils/formula'
+import { extractTitle } from '@renderer/utils/formats'
 import { Button } from 'antd'
 import { FC } from 'react'
 import { useTranslation } from 'react-i18next'

src/renderer/src/pages/home/Markdown/Markdown.tsx

@@ -2,7 +2,7 @@ import 'katex/dist/katex.min.css'
 
 import { useSettings } from '@renderer/hooks/useSettings'
 import { Message } from '@renderer/types'
-import { escapeBrackets, removeSvgEmptyLines } from '@renderer/utils/formula'
+import { escapeBrackets, removeSvgEmptyLines, withGeminiGrounding } from '@renderer/utils/formats'
 import { isEmpty } from 'lodash'
 import { FC, useMemo } from 'react'
 import { useTranslation } from 'react-i18next'
@@ -34,9 +34,9 @@ const Markdown: FC<Props> = ({ message }) => {
   const messageContent = useMemo(() => {
     const empty = isEmpty(message.content)
     const paused = message.status === 'paused'
-    const content = empty && paused ? t('message.chat.completion.paused') : message.content
+    const content = empty && paused ? t('message.chat.completion.paused') : withGeminiGrounding(message)
     return removeSvgEmptyLines(escapeBrackets(content))
-  }, [message.content, message.status, t])
+  }, [message, t])
 
   const rehypePlugins = useMemo(() => {
     const hasElements = ALLOWED_ELEMENTS.test(messageContent)

src/renderer/src/pages/home/Messages/MessageContent.tsx

@@ -10,6 +10,7 @@ import styled from 'styled-components'
 import Markdown from '../Markdown/Markdown'
 import MessageAttachments from './MessageAttachments'
 import MessageError from './MessageError'
+import MessageSearchResults from './MessageSearchResults'
 
 const MessageContent: React.FC<{
   message: Message
@@ -50,6 +51,7 @@ const MessageContent: React.FC<{
         </>
       )}
       <MessageAttachments message={message} />
+      <MessageSearchResults message={message} />
     </>
   )
 }

src/renderer/src/pages/home/Messages/MessageSearchResults.tsx

@@ -0,0 +1,96 @@
+import { InfoCircleOutlined } from '@ant-design/icons'
+import { Message } from '@renderer/types'
+import { FC } from 'react'
+import { useTranslation } from 'react-i18next'
+import styled from 'styled-components'
+
+interface Props {
+  message: Message
+}
+
+const MessageSearchResults: FC<Props> = ({ message }) => {
+  const { t } = useTranslation()
+
+  if (!message.metadata?.groundingMetadata) {
+    return null
+  }
+
+  const { groundingChunks, searchEntryPoint } = message.metadata.groundingMetadata
+
+  if (!groundingChunks) {
+    return null
+  }
+
+  let searchEntryContent = searchEntryPoint?.renderedContent
+
+  searchEntryContent = searchEntryContent?.replace(
+    /@media \(prefers-color-scheme: light\)/g,
+    'body[theme-mode="light"]'
+  )
+
+  searchEntryContent = searchEntryContent?.replace(/@media \(prefers-color-scheme: dark\)/g, 'body[theme-mode="dark"]')
+
+  return (
+    <>
+      <Container className="footnotes">
+        <TitleRow>
+          <Title>{t('common.footnotes')}</Title>
+          <InfoCircleOutlined />
+        </TitleRow>
+        <Sources>
+          {groundingChunks.map((chunk, index) => (
+            <SourceItem key={index}>
+              <Link href={chunk.web?.uri} target="_blank" rel="noopener noreferrer">
+                {chunk.web?.title}
+              </Link>
+            </SourceItem>
+          ))}
+        </Sources>
+      </Container>
+      <SearchEntryPoint dangerouslySetInnerHTML={{ __html: searchEntryContent || '' }} />
+    </>
+  )
+}
+
+const Container = styled.div`
+  padding: 16px;
+  border-radius: 8px;
+  margin-bottom: 0;
+`
+
+const TitleRow = styled.div`
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+  gap: 5px;
+  margin-bottom: 10px;
+`
+
+const Title = styled.h4`
+  margin: 0 !important;
+`
+
+const Sources = styled.ol`
+  margin-top: 10px;
+`
+
+const SourceItem = styled.li`
+  margin-bottom: 5px;
+`
+
+const Link = styled.a`
+  margin-left: 5px;
+  color: var(--color-primary);
+  text-decoration: none;
+
+  &:hover {
+    text-decoration: underline;
+  }
+`
+
+const SearchEntryPoint = styled.div`
+  margin-top: 10px;
+  margin-bottom: 10px;
+`
+
+export default MessageSearchResults

src/renderer/src/providers/GeminiProvider.ts

@@ -176,7 +176,8 @@ export default class GeminiProvider extends BaseProvider {
           completion_tokens: response.usageMetadata?.candidatesTokenCount,
           time_completion_millsec,
           time_first_token_millsec: 0
-        }
+        },
+        search: response.candidates?.[0]?.groundingMetadata
       })
       return
     }
@@ -201,7 +202,8 @@ export default class GeminiProvider extends BaseProvider {
           completion_tokens: chunk.usageMetadata?.candidatesTokenCount,
           time_completion_millsec,
           time_first_token_millsec
-        }
+        },
+        search: chunk.candidates?.[0]?.groundingMetadata
       })
     }
   }

src/renderer/src/providers/index.d.ts

@@ -1,14 +1,16 @@
+import type { GroundingMetadata } from '@google/generative-ai'
 import type { Assistant, Metrics } from '@renderer/types'
 
 interface ChunkCallbackData {
   text?: string
   usage?: OpenAI.Completions.CompletionUsage
   metrics?: Metrics
+  search?: GroundingMetadata
 }
 
 interface CompletionsParams {
   messages: Message[]
   assistant: Assistant
-  onChunk: ({ text, usage }: ChunkCallbackData) => void
+  onChunk: ({ text, usage, metrics, search }: ChunkCallbackData) => void
   onFilterMessages: (messages: Message[]) => void
 }

src/renderer/src/services/ApiService.ts

@@ -57,10 +57,15 @@ export async function fetchChatCompletion({
       messages,
       assistant,
       onFilterMessages: (messages) => (_messages = messages),
-      onChunk: ({ text, usage, metrics }) => {
+      onChunk: ({ text, usage, metrics, search }) => {
         message.content = message.content + text || ''
         message.usage = usage
         message.metrics = metrics
+
+        if (search) {
+          message.metadata = { groundingMetadata: search }
+        }
+
         onResponse({ ...message, status: 'pending' })
       }
     })

src/renderer/src/types/index.ts

@@ -59,6 +59,10 @@ export type Message = {
   knowledgeBaseIds?: string[]
   type: 'text' | '@' | 'clear'
   isPreset?: boolean
+  metadata?: {
+    // Gemini
+    groundingMetadata?: any
+  }
 }
 
 export type Metrics = {

src/renderer/src/utils/formats.ts

@@ -1,3 +1,5 @@
+import { Message } from '@renderer/types'
+
 export function escapeDollarNumber(text: string) {
   let escapedText = ''
 
@@ -56,3 +58,25 @@ export function removeSvgEmptyLines(text: string): string {
       .join('\n')
   })
 }
+
+export function withGeminiGrounding(message: Message) {
+  const { groundingSupports } = message?.metadata?.groundingMetadata || {}
+
+  if (!groundingSupports) {
+    return message.content
+  }
+
+  let content = message.content
+
+  groundingSupports.forEach((support) => {
+    const text = support.segment.text
+    const indices = support.groundingChunkIndices
+    const nodes = indices.reduce((acc, index) => {
+      acc.push(`<sup>${index + 1}</sup>`)
+      return acc
+    }, [])
+    content = content.replace(text, `${text} ${nodes.join(' ')}`)
+  })
+
+  return content
+}