fix: display all sources from annotations (#64)

* enhance UI of source nodes

* fix starter questions UI

* fix: display sources from all source anno

* fix: lint

* suggested questions not shown

* Create ten-dragons-shop.md

* fix lint

* remove log

.changeset/ten-dragons-shop.md

@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': patch
+---
+
+fix: display all sources from annotations

packages/chat-ui/src/chat/annotation.ts

@@ -130,15 +130,11 @@ export function getSourceAnnotationData(
     annotations,
     MessageAnnotationType.SOURCES
   )
-  if (data.length > 0) {
-    return [
-      {
-        ...data[0],
-        nodes: data[0].nodes ? preprocessSourceNodes(data[0].nodes) : [],
-      },
-    ]
-  }
-  return data
+  if (!data.length) return []
+  return data.map(item => ({
+    ...item,
+    nodes: item.nodes ? preprocessSourceNodes(item.nodes) : [],
+  }))
 }
 
 function preprocessSourceNodes(nodes: SourceNode[]): SourceNode[] {

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

@@ -18,6 +18,7 @@ import {
   SuggestedQuestionsData,
 } from './annotation'
 import { useChatMessage } from './chat-message.context.js'
+import { useChatUI } from './chat.context.js'
 
 export function EventAnnotations() {
   const { message, isLast, isLoading } = useChatMessage()
@@ -85,17 +86,17 @@ export function DocumentFileAnnotations() {
 export function SourceAnnotations() {
   const { message } = useChatMessage()
 
-  const annotations = message.annotations as MessageAnnotation[] | undefined
-  const sourceData =
-    annotations && annotations.length > 0
-      ? getSourceAnnotationData(annotations)
-      : null
-  if (!sourceData) return null
-  return sourceData[0] ? <ChatSources data={sourceData[0]} /> : null
+  const annotations = (message.annotations ?? []) as MessageAnnotation[]
+  const sourceData = getSourceAnnotationData(annotations)
+
+  if (!sourceData?.length) return null
+  const allNodes = sourceData.flatMap(item => item.nodes)
+  return <ChatSources data={{ nodes: allNodes }} />
 }
 
 export function SuggestedQuestionsAnnotations() {
-  const { message, append, isLast } = useChatMessage()
+  const { append } = useChatUI()
+  const { message, isLast } = useChatMessage()
   if (!isLast || !append) return null
 
   const annotations = message.annotations as MessageAnnotation[] | undefined

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

@@ -43,13 +43,13 @@ export function DocumentInfo({
   }
 
   const DocumentDetail = (
-    <div className={`relative ${className}`}>
+    <div className={`bg-secondary rounded-lg p-2 ${className}`}>
       <DocumentPreviewCard
         className="cursor-pointer"
         file={previewFile}
         onRemove={onRemove}
       />
-      <div className="absolute bottom-2 right-2 flex space-x-2">
+      <div className="flex max-w-60 flex-wrap space-x-2 px-2">
         {sources.map((node: SourceNode, index: number) => (
           <div key={node.id}>
             <SourceInfo node={node} index={startIndex + index} />

packages/chat-ui/src/widgets/starter-questions.tsx

@@ -9,12 +9,13 @@ interface StarterQuestionsProps {
 export function StarterQuestions(props: StarterQuestionsProps) {
   return (
     <div className="absolute bottom-6 left-0 w-full">
-      <div className="mx-20 grid grid-cols-2 gap-2">
+      <div className="mx-10 grid grid-cols-2 gap-2">
         {props.questions.map((question, i) => (
           <Button
             key={i}
             variant="outline"
             onClick={() => props.append({ role: 'user', content: question })}
+            className="h-auto whitespace-break-spaces"
           >
             {question}
           </Button>