fix: able to get custom annotation data (#42)

.changeset/sweet-fireants-relax.md

@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': patch
+---
+
+fix: able to get custom annotation data

.github/workflows/release.yml

@@ -0,0 +1,70 @@
+name: Release
+
+on:
+  push:
+    branches:
+      - main
+
+concurrency: ${{ github.workflow }}-${{ github.ref }}
+
+jobs:
+  release:
+    name: Release
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout Repo
+        uses: actions/checkout@v4
+
+      - uses: pnpm/action-setup@v4
+
+      - name: Setup Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version-file: ".nvmrc"
+          cache: "pnpm"
+
+      - name: Install dependencies
+        run: pnpm install
+
+      - name: Add auth token to .npmrc file
+        run: |
+          cat << EOF >> ".npmrc"
+            //registry.npmjs.org/:_authToken=$NPM_TOKEN
+          EOF
+        env:
+          NPM_TOKEN: ${{ secrets.NPM_TOKEN }}
+
+      - name: Get changeset status
+        id: get-changeset-status
+        run: |
+          pnpm changeset status --output .changeset/status.json
+          new_version=$(jq -r '.releases[] | select(.name == "llamaindex") | .newVersion' < .changeset/status.json)
+          rm -v .changeset/status.json
+          echo "new-version=${new_version}" >> "$GITHUB_OUTPUT"
+
+      - name: Create Release Pull Request or Publish to npm
+        id: changesets
+        uses: changesets/action@v1
+        with:
+          commit: Release ${{ steps.get-changeset-status.outputs.new-version }}
+          title: Release ${{ steps.get-changeset-status.outputs.new-version }}
+          # update version PR with the latest changesets
+          version: pnpm new-version
+          # build package and call changeset publish
+          publish: pnpm release
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          NPM_TOKEN: ${{ secrets.NPM_TOKEN }}
+
+      # Refs: https://github.com/changesets/changesets/issues/421
+      - name: Update lock file
+        continue-on-error: true
+        run: pnpm install --lockfile-only
+
+      - name: Commit lock file
+        continue-on-error: true
+        uses: stefanzweifel/git-auto-commit-action@v5
+        with:
+          commit_message: "chore: update lock file"
+          branch: changeset-release/main
+          file_pattern: "pnpm-lock.yaml"

apps/web/app/custom-demo.tsx

@@ -1,18 +1,20 @@
 'use client'
 
+import { Code } from '@/components/code'
+import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'
 import {
   ChatInput,
   ChatMessage,
   ChatMessages,
   ChatSection,
+  getCustomAnnotation,
+  JSONValue,
+  useChatMessage,
   useChatUI,
   useFile,
 } from '@llamaindex/chat-ui'
-import { Markdown } from '@llamaindex/chat-ui/widgets'
 import { Message, useChat } from 'ai/react'
 import { User2 } from 'lucide-react'
-import { Code } from '@/components/code'
-import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'
 
 const code = `
 import {
@@ -102,23 +104,22 @@ const initialMessages: Message[] = [
   },
 ]
 
-function Annotation({ annotations }: { annotations: any }) {
-  if (annotations && Array.isArray(annotations)) {
-    return annotations.map((annotation: any) => {
-      if (annotation.type === 'image' && annotation.url) {
-        return (
-          <img
-            alt="annotation"
-            className="h-[200px] object-contain"
-            key={annotation.url}
-            src={annotation.url}
-          />
-        )
-      }
-      return null
-    })
-  }
-  return null
+function Annotation() {
+  const { message } = useChatMessage()
+  const annotations = getCustomAnnotation<{ type: string; url: string }>(
+    message.annotations as JSONValue[],
+    a => a.type === 'image'
+  )
+
+  if (!annotations.length) return null
+  return annotations.map(annotation => (
+    <img
+      alt="annotation"
+      className="h-[200px] object-contain"
+      key={annotation.url}
+      src={annotation.url}
+    />
+  ))
 }
 
 function CustomChatMessagesList() {
@@ -144,8 +145,8 @@ function CustomChatMessagesList() {
             isLoading={isLoading}
             append={append}
           >
-            <Markdown content={message.content} />
-            <Annotation annotations={message.annotations} />
+            <Annotation />
+            <ChatMessage.Content.Markdown />
           </ChatMessage.Content>
           <ChatMessage.Actions />
         </ChatMessage>

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

@@ -1,3 +1,5 @@
+import { JSONValue } from './chat.interface'
+
 export enum MessageAnnotationType {
   IMAGE = 'image',
   DOCUMENT_FILE = 'document_file',
@@ -78,20 +80,53 @@ export type MessageAnnotation = {
 
 const NODE_SCORE_THRESHOLD = 0.25
 
-export function getAnnotationData<T extends AnnotationData>(
+/**
+ * Gets custom message annotations that don't match any standard MessageAnnotationType
+ * @param annotations - Array of message annotations to filter
+ * @param filter - Optional custom filter function to apply after filtering out standard annotations
+ * @returns Filtered array of custom message annotations
+ *
+ * First filters out any annotations that match MessageAnnotationType values,
+ * then applies the optional custom filter if provided.
+ */
+export function getCustomAnnotation<T = JSONValue>(
+  annotations: JSONValue[] | undefined,
+  filterFn?: (a: T) => boolean
+): T[] {
+  if (!Array.isArray(annotations) || !annotations.length) return [] as T[]
+  const customAnnotations = annotations.filter(
+    a => !isSupportedAnnotation(a)
+  ) as T[]
+  return filterFn ? customAnnotations.filter(filterFn) : customAnnotations
+}
+
+function isSupportedAnnotation(a: JSONValue): boolean {
+  return (
+    typeof a === 'object' &&
+    a !== null &&
+    a !== undefined &&
+    'type' in a &&
+    'data' in a &&
+    Object.values(MessageAnnotationType).includes(
+      a.type as MessageAnnotationType
+    )
+  )
+}
+
+export function getChatUIAnnotation<T extends AnnotationData>(
   annotations: MessageAnnotation[],
   type: string
 ): T[] {
   if (!annotations?.length) return []
   return annotations
-    .filter(a => a.type.toString() === type)
+    .filter(a => a && 'type' in a && a.type.toString() === type)
     .map(a => a.data as T)
 }
 
 export function getSourceAnnotationData(
   annotations: MessageAnnotation[]
 ): SourceData[] {
-  const data = getAnnotationData<SourceData>(
+  const data = getChatUIAnnotation<SourceData>(
     annotations,
     MessageAnnotationType.SOURCES
   )

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

@@ -10,36 +10,38 @@ import {
   AgentEventData,
   DocumentFileData,
   EventData,
-  getAnnotationData,
+  getChatUIAnnotation,
   getSourceAnnotationData,
   ImageData,
   MessageAnnotation,
   MessageAnnotationType,
   SuggestedQuestionsData,
 } from './annotation'
-import { ChatHandler, Message } from './chat.interface'
+import { useChatMessage } from './chat-message.context.js'
+
+export function EventAnnotations() {
+  const { message, isLast, isLoading } = useChatMessage()
+  const showLoading = (isLast && isLoading) ?? false
 
-export function EventAnnotations({
-  message,
-  showLoading,
-}: {
-  message: Message
-  showLoading: boolean
-}) {
   const annotations = message.annotations as MessageAnnotation[] | undefined
   const eventData =
     annotations && annotations.length > 0
-      ? getAnnotationData<EventData>(annotations, MessageAnnotationType.EVENTS)
+      ? getChatUIAnnotation<EventData>(
+          annotations,
+          MessageAnnotationType.EVENTS
+        )
       : null
   if (!eventData?.length) return null
   return <ChatEvents data={eventData} showLoading={showLoading} />
 }
 
-export function AgentEventAnnotations({ message }: { message: Message }) {
+export function AgentEventAnnotations() {
+  const { message } = useChatMessage()
+
   const annotations = message.annotations as MessageAnnotation[] | undefined
   const agentEventData =
     annotations && annotations.length > 0
-      ? getAnnotationData<AgentEventData>(
+      ? getChatUIAnnotation<AgentEventData>(
           annotations,
           MessageAnnotationType.AGENT_EVENTS
         )
@@ -53,21 +55,25 @@ export function AgentEventAnnotations({ message }: { message: Message }) {
   )
 }
 
-export function ImageAnnotations({ message }: { message: Message }) {
+export function ImageAnnotations() {
+  const { message } = useChatMessage()
+
   const annotations = message.annotations as MessageAnnotation[] | undefined
   const imageData =
     annotations && annotations.length > 0
-      ? getAnnotationData<ImageData>(annotations, 'image')
+      ? getChatUIAnnotation<ImageData>(annotations, 'image')
       : null
   if (!imageData) return null
   return imageData[0] ? <ChatImage data={imageData[0]} /> : null
 }
 
-export function DocumentFileAnnotations({ message }: { message: Message }) {
+export function DocumentFileAnnotations() {
+  const { message } = useChatMessage()
+
   const annotations = message.annotations as MessageAnnotation[] | undefined
   const contentFileData =
     annotations && annotations.length > 0
-      ? getAnnotationData<DocumentFileData>(
+      ? getChatUIAnnotation<DocumentFileData>(
           annotations,
           MessageAnnotationType.DOCUMENT_FILE
         )
@@ -76,7 +82,9 @@ export function DocumentFileAnnotations({ message }: { message: Message }) {
   return contentFileData[0] ? <ChatFiles data={contentFileData[0]} /> : null
 }
 
-export function SourceAnnotations({ message }: { message: Message }) {
+export function SourceAnnotations() {
+  const { message } = useChatMessage()
+
   const annotations = message.annotations as MessageAnnotation[] | undefined
   const sourceData =
     annotations && annotations.length > 0
@@ -86,17 +94,14 @@ export function SourceAnnotations({ message }: { message: Message }) {
   return sourceData[0] ? <ChatSources data={sourceData[0]} /> : null
 }
 
-export function SuggestedQuestionsAnnotations({
-  message,
-  append,
-}: {
-  message: Message
-  append: ChatHandler['append']
-}) {
+export function SuggestedQuestionsAnnotations() {
+  const { message, append, isLast } = useChatMessage()
+  if (!isLast || !append) return null
+
   const annotations = message.annotations as MessageAnnotation[] | undefined
   const suggestedQuestionsData =
     annotations && annotations.length > 0
-      ? getAnnotationData<SuggestedQuestionsData>(
+      ? getChatUIAnnotation<SuggestedQuestionsData>(
           annotations,
           MessageAnnotationType.SUGGESTED_QUESTIONS
         )

packages/chat-ui/src/chat/chat-message.context.ts

@@ -0,0 +1,20 @@
+import { createContext, useContext } from 'react'
+import { ChatHandler, Message } from './chat.interface'
+
+export interface ChatMessageContext {
+  message: Message
+  isLast: boolean
+  isLoading?: boolean
+  append?: ChatHandler['append']
+}
+
+export const chatMessageContext = createContext<ChatMessageContext | null>(null)
+
+export const ChatMessageProvider = chatMessageContext.Provider
+
+export const useChatMessage = () => {
+  const context = useContext(chatMessageContext)
+  if (!context)
+    throw new Error('useChatMessage must be used within a ChatMessageProvider')
+  return context
+}