Add Perplexity citations

src/renderer/components/Markdown.tsx

@@ -24,6 +24,7 @@ import ArrowForwardIosIcon from '@mui/icons-material/ArrowForwardIos'
 import ExpandLessIcon from '@mui/icons-material/ExpandLess'
 import PlayCircleIcon from '@mui/icons-material/PlayCircle'
 import NiceModal from '@ebay/nice-modal-react'
+import { MessageCitation } from '../../shared/types'
 
 export default function Markdown(props: {
   children: string
@@ -33,6 +34,7 @@ export default function Markdown(props: {
   className?: string
   generating?: boolean
   preferCollapsedCodeBlock?: boolean
+  citations?: MessageCitation[]
 }) {
   const {
     children,
@@ -42,7 +44,21 @@ export default function Markdown(props: {
     preferCollapsedCodeBlock,
     className,
     generating,
+    citations,
   } = props
+  // Append sources to the text content
+  const contentWithSources = useMemo(() => {
+    let content = children
+    if (citations && citations.length > 0) {
+      // Add sources section to the end of the content
+      content += '\n\nSources:\n'
+      citations.forEach((citation) => {
+        content += `[${citation.number}] ${citation.url}\n`
+      })
+    }
+
+    return content
+  }, [children, citations])
   return useMemo(
     () => (
       <ReactMarkdown
@@ -73,10 +89,18 @@ export default function Markdown(props: {
           ),
         }}
       >
-        {enableLaTeXRendering ? latex.processLaTeX(children) : children}
+        {enableLaTeXRendering ? latex.processLaTeX(contentWithSources) : contentWithSources}
       </ReactMarkdown>
     ),
-    [children, enableLaTeXRendering, enableMermaidRendering]
+    [
+      contentWithSources,
+      enableLaTeXRendering,
+      enableMermaidRendering,
+      hiddenCodeCopyButton,
+      generating,
+      preferCollapsedCodeBlock,
+      className,
+    ]
   )
 }
 

src/renderer/components/Message.tsx

@@ -560,6 +560,7 @@ const Message: FC<Props> = (props) => {
                                 autoCollapseCodeBlock &&
                                 (preferCollapsedCodeBlock || msg.role !== 'assistant' || previewArtifact)
                               }
+                              citations={msg.citations}
                             >
                               {item.text || ''}
                             </Markdown>

src/renderer/packages/models/abstract-ai-sdk.ts

@@ -150,6 +150,7 @@ export default abstract class AbstractAISDKModel implements ModelInterface {
     let contentParts: MessageContentParts = []
     let currentTextPart: MessageTextPart | undefined = undefined
     let reasoningContent = ''
+    const citations: Array<{ id: string; url: string; title: string }> = []
 
     for await (const chunk of result.fullStream) {
       console.debug('stream chunk', chunk)
@@ -188,6 +189,16 @@ export default abstract class AbstractAISDKModel implements ModelInterface {
         currentTextPart = undefined
         const storageKey = await saveImage('response', `data:${chunk.mimeType};base64,${chunk.base64}`)
         contentParts.push({ type: 'image', storageKey })
+      } else if (chunk.type === 'source') {
+        // Perplexity citations come as source chunks
+        const url = chunk.source?.url
+        if (url && !citations.find((c) => c.url === url)) {
+          citations.push({ 
+            id: chunk.source.id || url, 
+            url, 
+            title: chunk.source.title || url 
+          })
+        }
       } else if (chunk.type === 'error') {
         const error = chunk.error
         if (APICallError.isInstance(error)) {
@@ -204,14 +215,22 @@ export default abstract class AbstractAISDKModel implements ModelInterface {
     }
 
     const usage = await result.usage
+
+    // Process citations with numbers
+    const processedCitations = citations.map((citation, index) => ({
+      ...citation,
+      number: index + 1,
+    }))
+
     options.onResultChange?.({
       contentParts,
       reasoningContent,
       tokenCount: usage?.completionTokens,
       tokensUsed: usage?.totalTokens,
+      citations: processedCitations,
     })
 
-    return { contentParts, reasoningContent, usage }
+    return { contentParts, reasoningContent, usage, citations: processedCitations }
   }
 }
 

src/renderer/packages/models/types.ts

@@ -1,5 +1,5 @@
 import { ToolSet } from 'ai'
-import { Message, MessageContentParts, ProviderOptions, StreamTextResult } from 'src/shared/types'
+import { Message, MessageContentParts, MessageCitation, ProviderOptions, StreamTextResult } from 'src/shared/types'
 
 export interface ModelInterface {
   name: string
@@ -24,6 +24,7 @@ export interface ResultChange {
   contentParts?: MessageContentParts
   tokenCount?: number // 当前消息的 token 数量
   tokensUsed?: number // 生成当前消息的 token 使用量
+  citations?: MessageCitation[] // Perplexity AI citations
 }
 
 export type onResultChangeWithCancel = (data: ResultChange & { cancel?: () => void }) => void

src/shared/types.ts

@@ -29,6 +29,13 @@ export interface MessageLink {
   chatboxAILinkUUID?: string
 }
 
+export interface MessageCitation {
+  id: string
+  url: string
+  title: string
+  number: number // Citation number like [1], [2], etc.
+}
+
 export interface MessagePicture {
   url?: string
   storageKey?: string
@@ -60,6 +67,7 @@ export type StreamTextResult = {
   contentParts: MessageContentParts
   reasoningContent?: string
   usage?: LanguageModelUsage
+  citations?: MessageCitation[]
 }
 
 // Chatbox 应用的消息类型
@@ -82,6 +90,7 @@ export interface Message {
 
   files?: MessageFile[] // chatboxai 专用
   links?: MessageLink[] // chatboxai 专用
+  citations?: MessageCitation[] // Perplexity AI citations
 
   // webBrowsing?: MessageWebBrowsing // chatboxai 专用, （已废弃）
   // toolCalls?: MessageToolCalls // 已废弃，使用contentParts代替