diff --git a/.changeset/sweet-fireants-relax.md b/.changeset/sweet-fireants-relax.md
new file mode 100644
index 0000000..d090c08
--- /dev/null
+++ b/.changeset/sweet-fireants-relax.md
@@ -0,0 +1,5 @@
+---
+'@llamaindex/chat-ui': patch
+---
+
+fix: able to get custom annotation data
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
new file mode 100644
index 0000000..0cd5f69
--- /dev/null
+++ b/.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"
diff --git a/apps/web/app/custom-demo.tsx b/apps/web/app/custom-demo.tsx
index a6773b8..ffa88fe 100644
--- a/apps/web/app/custom-demo.tsx
+++ b/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>
diff --git a/packages/chat-ui/src/chat/annotation.ts b/packages/chat-ui/src/chat/annotation.ts
index 5091036..0575395 100644
--- a/packages/chat-ui/src/chat/annotation.ts
+++ b/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
   )
diff --git a/packages/chat-ui/src/chat/chat-annotations.tsx b/packages/chat-ui/src/chat/chat-annotations.tsx
index 49ca9e1..9f06680 100644
--- a/packages/chat-ui/src/chat/chat-annotations.tsx
+++ b/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
         )
diff --git a/packages/chat-ui/src/chat/chat-message.context.ts b/packages/chat-ui/src/chat/chat-message.context.ts
new file mode 100644
index 0000000..f78e645
--- /dev/null
+++ b/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
+}
diff --git a/packages/chat-ui/src/chat/chat-message.tsx b/packages/chat-ui/src/chat/chat-message.tsx
index b9ea57a..bcf4cb3 100644
--- a/packages/chat-ui/src/chat/chat-message.tsx
+++ b/packages/chat-ui/src/chat/chat-message.tsx
@@ -1,5 +1,5 @@
 import { Bot, Check, Copy, MessageCircle, User2 } from 'lucide-react'
-import { createContext, Fragment, memo, useContext, useMemo } from 'react'
+import { memo } from 'react'
 import { useCopyToClipboard } from '../hook/use-copy-to-clipboard'
 import { cn } from '../lib/utils'
 import { Button } from '../ui/button'
@@ -13,6 +13,7 @@ import {
   SourceAnnotations,
   SuggestedQuestionsAnnotations,
 } from './chat-annotations'
+import { ChatMessageProvider, useChatMessage } from './chat-message.context.js'
 import { ChatHandler, Message } from './chat.interface'
 
 interface ChatMessageProps extends React.PropsWithChildren {
@@ -47,14 +48,8 @@ export enum ContentPosition {
   BOTTOM = 9999,
 }
 
-type ContentDisplayConfig = {
-  position: ContentPosition
-  component: React.ReactNode | null
-}
-
 interface ChatMessageContentProps extends React.PropsWithChildren {
   className?: string
-  content?: ContentDisplayConfig[]
   isLoading?: boolean
   append?: ChatHandler['append']
   message?: Message // in case you want to customize the message
@@ -64,22 +59,6 @@ interface ChatMessageActionsProps extends React.PropsWithChildren {
   className?: string
 }
 
-interface ChatMessageContext {
-  message: Message
-  isLast: boolean
-}
-
-const chatMessageContext = createContext<ChatMessageContext | null>(null)
-
-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
-}
-
 function ChatMessage(props: ChatMessageProps) {
   const children = props.children ?? (
     <>
@@ -91,7 +70,12 @@ function ChatMessage(props: ChatMessageProps) {
 
   return (
     <ChatMessageProvider
-      value={{ message: props.message, isLast: props.isLast }}
+      value={{
+        message: props.message,
+        isLast: props.isLast,
+        isLoading: props.isLoading,
+        append: props.append,
+      }}
     >
       <div className={cn('group flex gap-4 p-3', props.className)}>
         {children}
@@ -120,73 +104,15 @@ function ChatMessageAvatar(props: ChatMessageAvatarProps) {
 }
 
 function ChatMessageContent(props: ChatMessageContentProps) {
-  const { message: defaultMessage, isLast } = useChatMessage()
-  const message = props.message ?? defaultMessage
-  const annotations = message.annotations as MessageAnnotation[] | undefined
-
-  const contents = useMemo<ContentDisplayConfig[]>(() => {
-    const displayMap: {
-      [key in ContentPosition]?: React.ReactNode | null
-    } = {
-      [ContentPosition.CHAT_EVENTS]: (
-        <EventAnnotations
-          message={message}
-          showLoading={(isLast && props.isLoading) ?? false}
-        />
-      ),
-      [ContentPosition.CHAT_AGENT_EVENTS]: (
-        <AgentEventAnnotations message={message} />
-      ),
-      [ContentPosition.CHAT_IMAGE]: <ImageAnnotations message={message} />,
-      [ContentPosition.MARKDOWN]: (
-        <Markdown
-          content={message.content}
-          sources={
-            annotations ? getSourceAnnotationData(annotations)[0] : undefined
-          }
-        />
-      ),
-      [ContentPosition.CHAT_DOCUMENT_FILES]: (
-        <DocumentFileAnnotations message={message} />
-      ),
-      [ContentPosition.CHAT_SOURCES]: <SourceAnnotations message={message} />,
-      ...(isLast &&
-        props.append && {
-          // show suggested questions only on the last message
-          [ContentPosition.SUGGESTED_QUESTIONS]: (
-            <SuggestedQuestionsAnnotations
-              message={message}
-              append={props.append}
-            />
-          ),
-        }),
-    }
-
-    // Override the default display map with the custom content
-    props.content?.forEach(content => {
-      displayMap[content.position] = content.component
-    })
-
-    return Object.entries(displayMap).map(([position, component]) => ({
-      position: parseInt(position),
-      component,
-    }))
-  }, [
-    annotations,
-    isLast,
-    message,
-    props.append,
-    props.content,
-    props.isLoading,
-  ])
-
   const children = props.children ?? (
     <>
-      {contents
-        .sort((a, b) => a.position - b.position)
-        .map((content, index) => (
-          <Fragment key={index}>{content.component}</Fragment>
-        ))}
+      <EventAnnotations />
+      <AgentEventAnnotations />
+      <ImageAnnotations />
+      <ChatMarkdown />
+      <DocumentFileAnnotations />
+      <SourceAnnotations />
+      <SuggestedQuestionsAnnotations />
     </>
   )
 
@@ -197,6 +123,20 @@ function ChatMessageContent(props: ChatMessageContentProps) {
   )
 }
 
+function ChatMarkdown() {
+  const { message } = useChatMessage()
+  const annotations = message.annotations as MessageAnnotation[] | undefined
+
+  return (
+    <Markdown
+      content={message.content}
+      sources={
+        annotations ? getSourceAnnotationData(annotations)[0] : undefined
+      }
+    />
+  )
+}
+
 function ChatMessageActions(props: ChatMessageActionsProps) {
   const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 })
   const { message } = useChatMessage()
@@ -223,9 +163,19 @@ function ChatMessageActions(props: ChatMessageActionsProps) {
   )
 }
 
+type ComposibleChatMessageContent = typeof ChatMessageContent & {
+  Event: typeof EventAnnotations
+  AgentEvent: typeof AgentEventAnnotations
+  Image: typeof ImageAnnotations
+  Markdown: typeof ChatMarkdown
+  DocumentFile: typeof DocumentFileAnnotations
+  Source: typeof SourceAnnotations
+  SuggestedQuestions: typeof SuggestedQuestionsAnnotations
+}
+
 type ComposibleChatMessage = typeof ChatMessage & {
   Avatar: typeof ChatMessageAvatar
-  Content: typeof ChatMessageContent
+  Content: ComposibleChatMessageContent
   Actions: typeof ChatMessageActions
 }
 
@@ -237,8 +187,18 @@ const PrimiviteChatMessage = memo(ChatMessage, (prevProps, nextProps) => {
   )
 }) as unknown as ComposibleChatMessage
 
+PrimiviteChatMessage.Content =
+  ChatMessageContent as ComposibleChatMessageContent
+
+PrimiviteChatMessage.Content.Event = EventAnnotations
+PrimiviteChatMessage.Content.AgentEvent = AgentEventAnnotations
+PrimiviteChatMessage.Content.Image = ImageAnnotations
+PrimiviteChatMessage.Content.Markdown = ChatMarkdown
+PrimiviteChatMessage.Content.DocumentFile = DocumentFileAnnotations
+PrimiviteChatMessage.Content.Source = SourceAnnotations
+PrimiviteChatMessage.Content.SuggestedQuestions = SuggestedQuestionsAnnotations
+
 PrimiviteChatMessage.Avatar = ChatMessageAvatar
-PrimiviteChatMessage.Content = ChatMessageContent
 PrimiviteChatMessage.Actions = ChatMessageActions
 
 export default PrimiviteChatMessage
diff --git a/packages/chat-ui/src/chat/chat.interface.ts b/packages/chat-ui/src/chat/chat.interface.ts
index 0197b3e..0d66b28 100644
--- a/packages/chat-ui/src/chat/chat.interface.ts
+++ b/packages/chat-ui/src/chat/chat.interface.ts
@@ -1,5 +1,15 @@
 type MessageRole = 'system' | 'user' | 'assistant' | 'data'
 
+export type JSONValue =
+  | null
+  | string
+  | number
+  | boolean
+  | {
+      [value: string]: JSONValue
+    }
+  | JSONValue[]
+
 export interface Message {
   content: string
   role: MessageRole
diff --git a/packages/chat-ui/src/index.tsx b/packages/chat-ui/src/index.tsx
index 71757e6..18b67d7 100644
--- a/packages/chat-ui/src/index.tsx
+++ b/packages/chat-ui/src/index.tsx
@@ -8,9 +8,9 @@ export { default as ChatMessage, ContentPosition } from './chat/chat-message'
 
 // Context Provider Hooks
 export { useChatUI } from './chat/chat.context'
+export { useChatMessage } from './chat/chat-message.context'
 export { useChatInput } from './chat/chat-input'
 export { useChatMessages } from './chat/chat-messages'
-export { useChatMessage } from './chat/chat-message'
 
 // Custom Hooks
 export { useFile } from './hook/use-file'