refactor: chat message content

templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx

-import {
-  ChatMessage,
-  ContentPosition,
-  getSourceAnnotationData,
-  useChatMessage,
-  useChatUI,
-} from "@llamaindex/chat-ui";
+import { ChatMessage } from "@llamaindex/chat-ui";
 import { DeepResearchCard } from "./custom/deep-research-card";
-import { Markdown } from "./custom/markdown";
 import { ToolAnnotations } from "./tools/chat-tools";
 
 export function ChatMessageContent() {
-  const { isLoading, append } = useChatUI();
-  const { message } = useChatMessage();
-  const customContent = [
-    {
-      // override the default markdown component
-      position: ContentPosition.MARKDOWN,
-      component: (
-        <Markdown
-          content={message.content}
-          sources={getSourceAnnotationData(message.annotations)?.[0]}
-        />
-      ),
-    },
-    // add the deep research card
-    {
-      position: ContentPosition.CHAT_EVENTS,
-      component: <DeepResearchCard message={message} />,
-    },
-    {
-      // add the tool annotations after events
-      position: ContentPosition.AFTER_EVENTS,
-      component: <ToolAnnotations message={message} />,
-    },
-  ];
   return (
-    <ChatMessage.Content
-      content={customContent}
-      isLoading={isLoading}
-      append={append}
-    />
+    <ChatMessage.Content>
+      <ChatMessage.Content.Event />
+      <ChatMessage.Content.AgentEvent />
+      <DeepResearchCard />
+      <ToolAnnotations />
+      <ChatMessage.Content.Image />
+      <ChatMessage.Content.Markdown />
+      <ChatMessage.Content.DocumentFile />
+      <ChatMessage.Content.Source />
+      <ChatMessage.Content.SuggestedQuestions />
+    </ChatMessage.Content>
   );
 }

templates/types/streaming/nextjs/app/components/ui/chat/custom/deep-research-card.tsx

 "use client";
 
-import { Message } from "@llamaindex/chat-ui";
+import { getCustomAnnotationData, useChatMessage } from "@llamaindex/chat-ui";
 import {
   AlertCircle,
   CheckCircle2,
@@ -54,7 +54,6 @@ type DeepResearchCardState = {
 };
 
 interface DeepResearchCardProps {
-  message: Message;
   className?: string;
 }
 
@@ -143,25 +142,19 @@ const deepResearchEventsToState = (
   );
 };
 
-export function DeepResearchCard({
-  message,
-  className,
-}: DeepResearchCardProps) {
-  const deepResearchEvents = message.annotations as
-    | DeepResearchEvent[]
-    | undefined;
-  const hasDeepResearchEvents = deepResearchEvents?.some(
-    (event) => event.type === "deep_research_event",
-  );
+export function DeepResearchCard({ className }: DeepResearchCardProps) {
+  const { message } = useChatMessage();
 
-  const state = useMemo(
-    () => deepResearchEventsToState(deepResearchEvents),
-    [deepResearchEvents],
-  );
+  const state = useMemo(() => {
+    const deepResearchEvents = getCustomAnnotationData<DeepResearchEvent>(
+      message.annotations,
+      (annotation) => annotation?.type === "deep_research_event",
+    );
+    if (!deepResearchEvents.length) return null;
+    return deepResearchEventsToState(deepResearchEvents);
+  }, [message.annotations]);
 
-  if (!hasDeepResearchEvents) {
-    return null;
-  }
+  if (!state) return null;
 
   return (
     <Card className={cn("w-full", className)}>

templates/types/streaming/nextjs/app/components/ui/chat/tools/chat-tools.tsx

@@ -2,6 +2,7 @@ import {
   Message,
   MessageAnnotation,
   getAnnotationData,
+  useChatMessage,
   useChatUI,
 } from "@llamaindex/chat-ui";
 import { JSONValue } from "ai";
@@ -9,10 +10,11 @@ import { useMemo } from "react";
 import { Artifact, CodeArtifact } from "./artifact";
 import { WeatherCard, WeatherData } from "./weather-card";
 
-export function ToolAnnotations({ message }: { message: Message }) {
+export function ToolAnnotations() {
   // TODO: This is a bit of a hack to get the artifact version. better to generate the version in the tool call and
   // store it in CodeArtifact
   const { messages } = useChatUI();
+  const { message } = useChatMessage();
   const artifactVersion = useMemo(
     () => getArtifactVersion(messages, message),
     [messages, message],