From 60eb08a982b4e2b62143d79375489738e8a4e046 Mon Sep 17 00:00:00 2001
From: shiquda <czy1847259360@gmail.com>
Date: Tue, 18 Mar 2025 16:48:07 +0800
Subject: [PATCH] feat: add support for citation preview (#3354)

* feat: add support for citation preview

#3217

* feat(MessageContent): Add HTML entity encoding to enhance the security of quoted data

* fix(MessageContent): recognize citation format like `[[1]]`
---
 .../pages/home/Markdown/CitationTooltip.tsx   | 109 ++++++++++++++++++
 src/renderer/src/pages/home/Markdown/Link.tsx |  46 +++++++-
 .../src/pages/home/Markdown/Markdown.tsx      |  19 ++-
 .../pages/home/Messages/MessageContent.tsx    |  56 ++++++++-
 4 files changed, 220 insertions(+), 10 deletions(-)
 create mode 100644 src/renderer/src/pages/home/Markdown/CitationTooltip.tsx

diff --git a/src/renderer/src/pages/home/Markdown/CitationTooltip.tsx b/src/renderer/src/pages/home/Markdown/CitationTooltip.tsx
new file mode 100644
index 00000000..b3d74ba7
--- /dev/null
+++ b/src/renderer/src/pages/home/Markdown/CitationTooltip.tsx
@@ -0,0 +1,109 @@
+import Favicon from '@renderer/components/Icons/FallbackFavicon'
+import { Tooltip } from 'antd'
+import React from 'react'
+import styled from 'styled-components'
+
+interface CitationTooltipProps {
+  children: React.ReactNode
+  citation: {
+    url: string
+    title?: string
+    content?: string
+  }
+}
+
+const CitationTooltip: React.FC<CitationTooltipProps> = ({ children, citation }) => {
+  let hostname = ''
+  try {
+    hostname = new URL(citation.url).hostname
+  } catch {
+    hostname = citation.url
+  }
+
+  // 自定义悬浮卡片内容
+  const tooltipContent = (
+    <TooltipContentWrapper>
+      <TooltipHeader onClick={() => window.open(citation.url, '_blank')}>
+        <Favicon hostname={hostname} alt={citation.title || hostname} />
+        <TooltipTitle title={citation.title || hostname}>{citation.title || hostname}</TooltipTitle>
+      </TooltipHeader>
+      {citation.content && <TooltipBody>{citation.content}</TooltipBody>}
+      <TooltipFooter onClick={() => window.open(citation.url, '_blank')}>{hostname}</TooltipFooter>
+    </TooltipContentWrapper>
+  )
+
+  return (
+    <StyledTooltip
+      title={tooltipContent}
+      placement="top"
+      arrow={false}
+      overlayInnerStyle={{
+        padding: 0,
+        borderRadius: '8px'
+      }}>
+      {children}
+    </StyledTooltip>
+  )
+}
+
+// 使用styled-components来自定义Tooltip的样式，包括箭头
+const StyledTooltip = styled(Tooltip)`
+  .ant-tooltip-arrow {
+    .ant-tooltip-arrow-content {
+      background-color: var(--color-background-1);
+    }
+  }
+`
+
+const TooltipContentWrapper = styled.div`
+  padding: 12px;
+  background-color: var(--color-background-soft);
+  border-radius: 8px;
+`
+
+const TooltipHeader = styled.div`
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  margin-bottom: 8px;
+  cursor: pointer;
+
+  &:hover {
+    opacity: 0.8;
+  }
+`
+
+const TooltipTitle = styled.div`
+  color: var(--color-text-1);
+  font-size: 14px;
+  line-height: 1.4;
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+`
+
+const TooltipBody = styled.div`
+  font-size: 13px;
+  line-height: 1.5;
+  margin-bottom: 8px;
+  overflow: hidden;
+  display: -webkit-box;
+  -webkit-line-clamp: 3;
+  -webkit-box-orient: vertical;
+  color: var(--color-text-2);
+`
+
+const TooltipFooter = styled.div`
+  font-size: 12px;
+  color: var(--color-link);
+  overflow: hidden;
+  text-overflow: ellipsis;
+  white-space: nowrap;
+  cursor: pointer;
+
+  &:hover {
+    text-decoration: underline;
+  }
+`
+
+export default CitationTooltip
diff --git a/src/renderer/src/pages/home/Markdown/Link.tsx b/src/renderer/src/pages/home/Markdown/Link.tsx
index 25ebc170..ce158035 100644
--- a/src/renderer/src/pages/home/Markdown/Link.tsx
+++ b/src/renderer/src/pages/home/Markdown/Link.tsx
@@ -1,12 +1,54 @@
 import { omit } from 'lodash'
 import React from 'react'
 
-const Link: React.FC = (props: React.AnchorHTMLAttributes<HTMLAnchorElement>) => {
+import CitationTooltip from './CitationTooltip'
+
+interface LinkProps extends React.AnchorHTMLAttributes<HTMLAnchorElement> {
+  node?: any
+  citationData?: {
+    url: string
+    title?: string
+    content?: string
+  }
+}
+
+const Link: React.FC<LinkProps> = (props) => {
+  // 处理内部链接
   if (props.href?.startsWith('#')) {
     return <span className="link">{props.children}</span>
   }
 
-  return <a {...omit(props, 'node')} target="_blank" rel="noreferrer" onClick={(e) => e.stopPropagation()} />
+  // 包含<sup>标签表示是一个引用链接
+  const isCitation = React.Children.toArray(props.children).some((child) => {
+    if (typeof child === 'object' && 'type' in child) {
+      return child.type === 'sup'
+    }
+    return false
+  })
+
+  // 如果是引用链接并且有引用数据，则使用CitationTooltip
+  if (isCitation && props.citationData) {
+    return (
+      <CitationTooltip citation={props.citationData}>
+        <a
+          {...omit(props, ['node', 'citationData'])}
+          target="_blank"
+          rel="noreferrer"
+          onClick={(e) => e.stopPropagation()}
+        />
+      </CitationTooltip>
+    )
+  }
+
+  // 普通链接
+  return (
+    <a
+      {...omit(props, ['node', 'citationData'])}
+      target="_blank"
+      rel="noreferrer"
+      onClick={(e) => e.stopPropagation()}
+    />
+  )
 }
 
 export default Link
diff --git a/src/renderer/src/pages/home/Markdown/Markdown.tsx b/src/renderer/src/pages/home/Markdown/Markdown.tsx
index 7d6306cc..5f255ced 100644
--- a/src/renderer/src/pages/home/Markdown/Markdown.tsx
+++ b/src/renderer/src/pages/home/Markdown/Markdown.tsx
@@ -26,9 +26,17 @@ const ALLOWED_ELEMENTS =
 
 interface Props {
   message: Message
+  citationsData?: Map<
+    string,
+    {
+      url: string
+      title?: string
+      content?: string
+    }
+  >
 }
 
-const Markdown: FC<Props> = ({ message }) => {
+const Markdown: FC<Props> = ({ message, citationsData }) => {
   const { t } = useTranslation()
   const { renderInputMessageAsMarkdown, mathEngine } = useSettings()
 
@@ -48,7 +56,12 @@ const Markdown: FC<Props> = ({ message }) => {
 
   const components = useCallback(() => {
     const baseComponents = {
-      a: Link,
+      a: (props: any) => {
+        if (props.href && citationsData?.has(props.href)) {
+          return <Link {...props} citationData={citationsData.get(props.href)} />
+        }
+        return <Link {...props} />
+      },
       code: CodeBlock,
       img: ImagePreview
     } as Partial<Components>
@@ -58,7 +71,7 @@ const Markdown: FC<Props> = ({ message }) => {
     }
 
     return baseComponents
-  }, [messageContent])
+  }, [messageContent, citationsData])
 
   if (message.role === 'user' && !renderInputMessageAsMarkdown) {
     return <p style={{ marginBottom: 5, whiteSpace: 'pre-wrap' }}>{messageContent}</p>
diff --git a/src/renderer/src/pages/home/Messages/MessageContent.tsx b/src/renderer/src/pages/home/Messages/MessageContent.tsx
index 02958b4f..6d5f7b69 100644
--- a/src/renderer/src/pages/home/Messages/MessageContent.tsx
+++ b/src/renderer/src/pages/home/Messages/MessageContent.tsx
@@ -29,6 +29,50 @@ const MessageContent: React.FC<Props> = ({ message: _message, model }) => {
   const { t } = useTranslation()
   const message = withMessageThought(clone(_message))
 
+  // HTML实体编码辅助函数
+  const encodeHTML = (str: string) => {
+    return str.replace(/[&<>"']/g, (match) => {
+      const entities: { [key: string]: string } = {
+        '&': '&amp;',
+        '<': '&lt;',
+        '>': '&gt;',
+        '"': '&quot;',
+        "'": '&apos;'
+      }
+      return entities[match]
+    })
+  }
+
+  // 获取引用数据
+  const citationsData = useMemo(() => {
+    const searchResults = message?.metadata?.webSearch?.results || []
+    const citationsUrls = message?.metadata?.citations || []
+
+    // 合并引用数据
+    const data = new Map()
+
+    // 添加webSearch结果
+    searchResults.forEach((result) => {
+      data.set(result.url, {
+        url: result.url,
+        title: result.title,
+        content: result.content
+      })
+    })
+
+    // 添加citations
+    citationsUrls.forEach((url) => {
+      if (!data.has(url)) {
+        data.set(url, {
+          url: url
+          // 如果没有title和content，将在CitationTooltip中显示hostname
+        })
+      }
+    })
+
+    return data
+  }, [message.metadata?.citations, message.metadata?.webSearch?.results])
+
   // Process content to make citation numbers clickable
   const processedContent = useMemo(() => {
     if (!(message.metadata?.citations || message.metadata?.webSearch)) {
@@ -42,18 +86,20 @@ const MessageContent: React.FC<Props> = ({ message: _message, model }) => {
     const citations = message?.metadata?.citations || searchResultsCitations
 
     // Convert [n] format to superscript numbers and make them clickable
-    // Use <sup> tag for superscript and make it a link
-    content = content.replace(/\[(\d+)\]/g, (match, num) => {
+    // Use <sup> tag for superscript and make it a link with citation data
+    content = content.replace(/\[\[(\d+)\]\]|\[(\d+)\]/g, (match, num1, num2) => {
+      const num = num1 || num2
       const index = parseInt(num) - 1
       if (index >= 0 && index < citations.length) {
         const link = citations[index]
-        return link ? `[<sup>${num}</sup>](${link})` : `<sup>${num}</sup>`
+        const citationData = link ? encodeHTML(JSON.stringify(citationsData.get(link) || { url: link })) : null
+        return link ? `[<sup data-citation='${citationData}'>${num}</sup>](${link})` : `<sup>${num}</sup>`
       }
       return match
     })
 
     return content
-  }, [message.content, message.metadata])
+  }, [message.content, message.metadata, citationsData])
 
   // Format citations for display
   const formattedCitations = useMemo(() => {
@@ -103,7 +149,7 @@ const MessageContent: React.FC<Props> = ({ message: _message, model }) => {
       </Flex>
       <MessageThought message={message} />
       <MessageTools message={message} />
-      <Markdown message={{ ...message, content: processedContent }} />
+      <Markdown message={{ ...message, content: processedContent }} citationsData={citationsData} />
       {message.translatedContent && (
         <Fragment>
           <Divider style={{ margin: 0, marginBottom: 10 }}>