fix: document editor with citations (#120)

* fix: document editor with citations

* fix: format

* process document editor content with citation numbers

* Create beige-wolves-watch.md

* fix format

* show node number

.changeset/beige-wolves-watch.md

@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': patch
+---
+
+feat: document editor with citations

packages/chat-ui/src/chat/canvas/artifacts/document.tsx

@@ -2,23 +2,41 @@
 
 import { FileText } from 'lucide-react'
 import { cn } from '../../../lib/utils'
-import { DocumentEditor } from '../../../widgets'
+import { DocumentEditor, SourceNode } from '../../../widgets'
 import { DocumentArtifact } from '../artifacts'
 import { ChatCanvasActions } from '../actions'
 import { useChatCanvas } from '../context'
 import { useState } from 'react'
 import { Button } from '../../../ui/button'
+import { useChatUI } from '../../chat.context'
+import { getSourceNodes } from '../../chat-annotations'
 
 interface DocumentArtifactViewerProps {
   className?: string
   children?: React.ReactNode
 }
 
+// Convert citation with node_id to citation number in markdown
+function processDocument(content: string, nodes: SourceNode[]) {
+  if (nodes.length === 0) return content
+
+  const citationRegex =
+    /\[citation:([a-fA-F0-9\\-]+)\]\(javascript:void\(0\)\)/g
+
+  return content.replace(citationRegex, (match, citationId) => {
+    const nodeIndex = nodes.findIndex(node => node.id === citationId)
+    if (nodeIndex !== -1) return ` \`${nodeIndex + 1}\` `
+    return match // return original citation if not found
+  })
+}
+
 export function DocumentArtifactViewer({
   className,
   children,
 }: DocumentArtifactViewerProps) {
   const { displayedArtifact, updateArtifact } = useChatCanvas()
+  const { messages } = useChatUI()
+
   const [updatedContent, setUpdatedContent] = useState<string | undefined>()
 
   if (displayedArtifact?.type !== 'document') return null
@@ -28,6 +46,9 @@ export function DocumentArtifactViewer({
     data: { content, title, type },
   } = documentArtifact
 
+  const nodes = messages.flatMap(message => getSourceNodes(message))
+  const transformedContent = processDocument(content, nodes)
+
   const handleDocumentChange = (markdown: string) => {
     setUpdatedContent(markdown)
   }
@@ -50,9 +71,9 @@ export function DocumentArtifactViewer({
         </h3>
         <ChatCanvasActions />
       </div>
-      <div className="relative mx-20 flex min-h-0 flex-1 flex-col items-stretch gap-4 py-4">
+      <div className="relative mx-20 flex min-h-0 flex-1 flex-col items-stretch gap-4 py-2">
         {updatedContent && (
-          <div className="bg-background absolute right-0 top-2 flex gap-2 py-2 pr-2 text-sm">
+          <div className="absolute right-[30px] top-[14px] z-20 flex gap-2 text-sm">
             <Button
               size="sm"
               className="h-7 bg-blue-500 hover:bg-blue-600"
@@ -73,8 +94,9 @@ export function DocumentArtifactViewer({
         {children ?? (
           <DocumentEditor
             key={documentArtifact.created_at}
-            content={content}
+            content={transformedContent}
             onChange={handleDocumentChange}
+            className="h-full overflow-y-auto"
           />
         )}
       </div>

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

@@ -1,4 +1,3 @@
-import { ComponentType } from 'react'
 import {
   ChatAgentEvents,
   ChatEvents,
@@ -19,7 +18,6 @@ import { getAnnotationData } from './annotations/annotations.js'
 import { useChatMessage } from './chat-message.context.js'
 import { useChatUI } from './chat.context.js'
 import { Message } from './chat.interface.js'
-import ChatCanvas from './canvas/index.js'
 
 export function EventAnnotations() {
   const { message, isLast, isLoading } = useChatMessage()
@@ -120,10 +118,3 @@ export function SuggestedQuestionsAnnotations() {
     />
   )
 }
-
-export const defaultAnnotationRenderers: Record<
-  string,
-  ComponentType<{ data: any }>
-> = {
-  artifact: ChatCanvas.Artifact,
-}

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

@@ -10,7 +10,6 @@ import {
 } from '../widgets/index.js'
 import {
   AgentEventAnnotations,
-  defaultAnnotationRenderers,
   DocumentFileAnnotations,
   EventAnnotations,
   ImageAnnotations,
@@ -21,6 +20,7 @@ import {
 import { ChatMessageProvider, useChatMessage } from './chat-message.context.js'
 import { useChatUI } from './chat.context.js'
 import { ChatHandler, Message } from './chat.interface'
+import { defaultAnnotationRenderers } from './chat-renderers.js'
 
 interface ChatMessageProps extends React.PropsWithChildren {
   message: Message

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

@@ -0,0 +1,9 @@
+import { ComponentType } from 'react'
+import ChatCanvas from './canvas/index.js'
+
+export const defaultAnnotationRenderers: Record<
+  string,
+  ComponentType<{ data: any }>
+> = {
+  artifact: ChatCanvas.Artifact,
+}

packages/chat-ui/src/index.tsx

@@ -7,7 +7,7 @@ export { default as ChatCanvas } from './chat/canvas'
 export { default as ChatInput } from './chat/chat-input'
 export { default as ChatMessages } from './chat/chat-messages'
 export { default as ChatMessage, ContentPosition } from './chat/chat-message'
-export { defaultAnnotationRenderers } from './chat/chat-annotations'
+export { defaultAnnotationRenderers } from './chat/chat-renderers'
 
 // Context Provider Hooks
 export { useChatUI } from './chat/chat.context'

packages/chat-ui/src/widgets/document-editor.tsx

@@ -10,6 +10,8 @@ import {
   markdownShortcutPlugin,
   tablePlugin,
   toolbarPlugin,
+  thematicBreakPlugin,
+  quotePlugin,
 } from '@mdxeditor/editor'
 
 export function DocumentEditor({
@@ -30,6 +32,8 @@ export function DocumentEditor({
     linkDialogPlugin(),
     tablePlugin(),
     markdownShortcutPlugin(),
+    thematicBreakPlugin(),
+    quotePlugin(),
   ]
 
   if (showToolbar) {
@@ -53,6 +57,12 @@ export function DocumentEditor({
       markdown={content}
       plugins={plugins}
       contentEditableClassName="custom-markdown"
+      onError={error => {
+        console.warn(
+          '[Chat-UI] Error while parsing markdown in DocumentEditor',
+          error
+        )
+      }}
     />
   )
 }