Fix not showing citation issue (#56)

.changeset/calm-bananas-smile.md

@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': minor
+---
+
+Fix not showing citation issue

packages/chat-ui/src/chat/chat-message.tsx

@@ -1,9 +1,9 @@
 import { Bot, Check, Copy, MessageCircle, User2 } from 'lucide-react'
-import { memo } from 'react'
+import { memo, ComponentType } from 'react'
 import { useCopyToClipboard } from '../hook/use-copy-to-clipboard'
 import { cn } from '../lib/utils'
 import { Button } from '../ui/button'
-import { Markdown } from '../widgets/index.js' // this import needs the file extension as it's importing the widget bundle
+import { Markdown, CitationComponentProps } from '../widgets/index.js'
 import { getSourceAnnotationData, MessageAnnotation } from './annotation'
 import {
   AgentEventAnnotations,
@@ -59,6 +59,10 @@ interface ChatMessageActionsProps extends React.PropsWithChildren {
   className?: string
 }
 
+interface ChatMarkdownProps extends React.PropsWithChildren {
+  citationComponent?: ComponentType<CitationComponentProps>
+}
+
 function ChatMessage(props: ChatMessageProps) {
   const children = props.children ?? (
     <>
@@ -123,7 +127,7 @@ function ChatMessageContent(props: ChatMessageContentProps) {
   )
 }
 
-function ChatMarkdown() {
+function ChatMarkdown(props: ChatMarkdownProps) {
   const { message } = useChatMessage()
   const annotations = message.annotations as MessageAnnotation[] | undefined
 
@@ -133,6 +137,7 @@ function ChatMarkdown() {
       sources={
         annotations ? getSourceAnnotationData(annotations)[0] : undefined
       }
+      citationComponent={props.citationComponent}
     />
   )
 }

packages/chat-ui/src/widgets/citation.tsx

@@ -0,0 +1,11 @@
+import { SourceNode } from '../chat/annotation'
+import { SourceNumberButton } from './source-number-button'
+
+export interface CitationComponentProps {
+  index: number
+  node: SourceNode
+}
+
+export function Citation({ index }: CitationComponentProps) {
+  return <SourceNumberButton index={index} />
+}

packages/chat-ui/src/widgets/index.tsx

@@ -12,3 +12,4 @@ export { StarterQuestions } from './starter-questions'
 export { DocumentInfo } from './document-info'
 export { ImagePreview } from './image-preview'
 export { FileUploader } from './file-uploader'
+export { Citation, type CitationComponentProps } from './citation'

packages/chat-ui/src/widgets/markdown.tsx

@@ -1,4 +1,4 @@
-import { FC, memo } from 'react'
+import { FC, memo, ComponentType } from 'react'
 import ReactMarkdown, { Options } from 'react-markdown'
 import rehypeKatex from 'rehype-katex'
 import remarkGfm from 'remark-gfm'
@@ -10,7 +10,7 @@ import {
   SourceData,
 } from '../chat/annotation'
 import { DocumentInfo } from './document-info'
-import { SourceNumberButton } from './source-number-button'
+import { Citation, CitationComponentProps } from './citation'
 
 const MemoizedReactMarkdown: FC<Options> = memo(
   ReactMarkdown,
@@ -40,30 +40,27 @@ const preprocessMedia = (content: string) => {
 }
 
 /**
- * Update the citation flag [citation:id]() to the new format [citation:index](url)
+ * Convert citation flags [citation:id] to markdown links [citation:id]()
  */
 const preprocessCitations = (input: string, sources?: SourceData) => {
   let content = input
+
   if (sources) {
-    const citationRegex = /\[citation:(.+?)\]/g
-    let match
-    // Find all the citation references in the content
-    while ((match = citationRegex.exec(content)) !== null) {
-      const citationId = match[1]
-      // Find the source node with the id equal to the citation-id, also get the index of the source node
-      const sourceNode = sources.nodes.find(node => node.id === citationId)
-      // If the source node is found, replace the citation reference with the new format
-      if (sourceNode !== undefined) {
-        content = content.replace(
-          match[0],
-          `[citation:${sources.nodes.indexOf(sourceNode)}]()`
-        )
-      } else {
-        // If the source node is not found, remove the citation reference
-        content = content.replace(match[0], '')
-      }
-    }
+    // Match citation format [citation:node_id]
+    // Handle complete citations
+    const idToIndexRegex = /\[citation:([^\]]+)\]/g
+    content = content.replace(idToIndexRegex, (match, citationId) => {
+      const trimmedId = citationId.trim()
+      // Use a special format that doesn't get styled as a link by markdown-it
+      return `[citation:${trimmedId}](javascript:void(0))`
+    })
+
+    // For incomplete citations - any [citation: pattern that isn't closed with ]
+    // Look for open bracket, citation text, then end of string or any char except closing bracket
+    const incompleteRegex = /\[citation:[^\]]*$/g
+    content = content.replace(incompleteRegex, '')
   }
+
   return content
 }
 
@@ -75,10 +72,12 @@ export function Markdown({
   content,
   sources,
   backend,
+  citationComponent: CitationComponent,
 }: {
   content: string
   sources?: SourceData
   backend?: string
+  citationComponent?: ComponentType<CitationComponentProps>
 }) {
   const processedContent = preprocessContent(content, sources)
 
@@ -146,18 +145,33 @@ export function Markdown({
                 )
               }
             }
-            // If a text link starts with 'citation:', then render it as a citation reference
+
+            // Handle citation links
             if (
               Array.isArray(children) &&
               typeof children[0] === 'string' &&
-              children[0].startsWith('citation:')
+              (children[0].startsWith('citation:') ||
+                href?.startsWith('citation:'))
             ) {
-              const index = Number(children[0].replace('citation:', ''))
-              if (!isNaN(index)) {
-                return <SourceNumberButton index={index} />
+              // Extract the nodeId from the citation link
+              const nodeId = children[0].includes('citation:')
+                ? children[0].split('citation:')[1].trim()
+                : href?.replace('citation:', '').trim() || ''
+
+              const nodeIndex = sources?.nodes.findIndex(
+                node => node.id === nodeId
+              )
+              const sourceNode = sources?.nodes.find(node => node.id === nodeId)
+
+              if (nodeIndex !== undefined && nodeIndex > -1 && sourceNode) {
+                return CitationComponent ? (
+                  <CitationComponent index={nodeIndex} node={sourceNode} />
+                ) : (
+                  <Citation index={nodeIndex} node={sourceNode} />
+                )
+              } else {
+                return null
               }
-              // citation is not looked up yet, don't render anything
-              return null
             }
             return (
               <a href={href} target="_blank" rel="noopener">

packages/chat-ui/src/widgets/source-number-button.tsx

@@ -10,7 +10,7 @@ export function SourceNumberButton({
   return (
     <span
       className={cn(
-        'inline-flex h-5 w-5 items-center justify-center rounded-full bg-gray-100 text-xs',
+        'inline-flex h-5 w-5 items-center justify-center rounded-full bg-gray-100 text-xs text-black',
         className
       )}
     >