Resolves #378

Author: haydenbleasel

.changeset/funny-bikes-own.md

@@ -0,0 +1,5 @@
+---
+"ai-elements": patch
+---
+
+Update ConversationMessage to UIMessage

apps/docs/content/components/(chatbot)/conversation.mdx

@@ -250,7 +250,7 @@ import { ConversationDownload } from "@/components/ai-elements/conversation";
   type={{
     messages: {
       description: "Array of messages to include in the download.",
-      type: "ConversationMessage[]",
+      type: "UIMessage[]",
       required: true,
     },
     filename: {
@@ -260,7 +260,7 @@ import { ConversationDownload } from "@/components/ai-elements/conversation";
     },
     formatMessage: {
       description: "Custom function to format each message in the output.",
-      type: "(message: ConversationMessage, index: number) => string",
+      type: "(message: UIMessage, index: number) => string",
     },
     "...props": {
       description:
@@ -282,6 +282,6 @@ const markdown = messagesToMarkdown(messages);
 // With custom formatter
 const customMarkdown = messagesToMarkdown(
   messages,
-  (msg, i) => `[${msg.role}]: ${msg.content}`
+  (msg, i) => `[${msg.role}]: ${msg.parts.filter(p => p.type === "text").map(p => p.text).join("")}`
 );
 ```

packages/elements/__tests__/conversation.test.tsx

@@ -66,8 +66,16 @@ vi.mock<typeof import("use-stick-to-bottom")>(
 );
 
 // Custom format function for messagesToMarkdown test
-const customFormatMessage = (msg: { role: string; content: string }) =>
-  `[${msg.role}]: ${msg.content}`;
+const customFormatMessage = (msg: {
+  role: string;
+  parts: { type: string; text?: string }[];
+}) => {
+  const text = msg.parts
+    .filter((p) => p.type === "text")
+    .map((p) => p.text)
+    .join("");
+  return `[${msg.role}]: ${text}`;
+};
 
 describe("conversation", () => {
   it("renders children", () => {
@@ -219,13 +227,19 @@ describe("conversationScrollButton", () => {
   });
 });
 
+const makeMessage = (role: "user" | "assistant" | "system", text: string) => ({
+  id: `${role}-${text}`,
+  parts: [{ text, type: "text" as const }],
+  role,
+});
+
 // Function name as describe title is a valid testing pattern
 // oxlint-disable-next-line eslint-plugin-jest(valid-title)
 describe(messagesToMarkdown, () => {
   it("converts messages to markdown format", () => {
     const messages = [
-      { content: "Hello", role: "user" as const },
-      { content: "Hi there!", role: "assistant" as const },
+      makeMessage("user", "Hello"),
+      makeMessage("assistant", "Hi there!"),
     ];
 
     const result = messagesToMarkdown(messages);
@@ -240,8 +254,8 @@ describe(messagesToMarkdown, () => {
 
   it("uses custom formatMessage function", () => {
     const messages = [
-      { content: "Hello", role: "user" as const },
-      { content: "Hi", role: "assistant" as const },
+      makeMessage("user", "Hello"),
+      makeMessage("assistant", "Hi"),
     ];
 
     const result = messagesToMarkdown(messages, customFormatMessage);
@@ -251,20 +265,39 @@ describe(messagesToMarkdown, () => {
 
   it("handles all role types", () => {
     const messages = [
-      { content: "User msg", role: "user" as const },
-      { content: "Assistant msg", role: "assistant" as const },
-      { content: "System msg", role: "system" as const },
-      { content: "Tool msg", role: "tool" as const },
-      { content: "Data msg", role: "data" as const },
+      makeMessage("user", "User msg"),
+      makeMessage("assistant", "Assistant msg"),
+      makeMessage("system", "System msg"),
     ];
 
     const result = messagesToMarkdown(messages);
 
     expect(result).toContain("**User:** User msg");
     expect(result).toContain("**Assistant:** Assistant msg");
     expect(result).toContain("**System:** System msg");
-    expect(result).toContain("**Tool:** Tool msg");
-    expect(result).toContain("**Data:** Data msg");
+  });
+
+  it("extracts text from multiple parts", () => {
+    const message = {
+      id: "multi",
+      parts: [
+        { text: "Hello ", type: "text" as const },
+        {
+          args: {},
+          result: {},
+          state: "result" as const,
+          toolInvocationId: "1",
+          toolName: "test",
+          type: "tool-invocation" as const,
+        },
+        { text: "world", type: "text" as const },
+      ],
+      role: "assistant" as const,
+    };
+
+    const result = messagesToMarkdown([message]);
+
+    expect(result).toBe("**Assistant:** Hello world");
   });
 });
 
@@ -310,8 +343,8 @@ const setupDomClickTracker = () => {
 
 describe("conversationDownload", () => {
   const mockMessages = [
-    { content: "Hello", role: "user" as const },
-    { content: "Hi there!", role: "assistant" as const },
+    makeMessage("user", "Hello"),
+    makeMessage("assistant", "Hi there!"),
   ];
 
   it("renders download button", () => {

packages/elements/src/conversation.tsx

@@ -2,6 +2,7 @@
 
 import { Button } from "@repo/shadcn-ui/components/ui/button";
 import { cn } from "@repo/shadcn-ui/lib/utils";
+import type { UIMessage } from "ai";
 import { ArrowDownIcon, DownloadIcon } from "lucide-react";
 import type { ComponentProps } from "react";
 import { useCallback } from "react";
@@ -99,30 +100,31 @@ export const ConversationScrollButton = ({
   );
 };
 
-export interface ConversationMessage {
-  role: "user" | "assistant" | "system" | "data" | "tool";
-  content: string;
-}
+const getMessageText = (message: UIMessage): string =>
+  message.parts
+    .filter((part) => part.type === "text")
+    .map((part) => part.text)
+    .join("");
 
 export type ConversationDownloadProps = Omit<
   ComponentProps<typeof Button>,
   "onClick"
 > & {
-  messages: ConversationMessage[];
+  messages: UIMessage[];
   filename?: string;
-  formatMessage?: (message: ConversationMessage, index: number) => string;
+  formatMessage?: (message: UIMessage, index: number) => string;
 };
 
-const defaultFormatMessage = (message: ConversationMessage): string => {
+const defaultFormatMessage = (message: UIMessage): string => {
   const roleLabel =
     message.role.charAt(0).toUpperCase() + message.role.slice(1);
-  return `**${roleLabel}:** ${message.content}`;
+  return `**${roleLabel}:** ${getMessageText(message)}`;
 };
 
 export const messagesToMarkdown = (
-  messages: ConversationMessage[],
+  messages: UIMessage[],
   formatMessage: (
-    message: ConversationMessage,
+    message: UIMessage,
     index: number
   ) => string = defaultFormatMessage
 ): string => messages.map((msg, i) => formatMessage(msg, i)).join("\n\n");

skills/ai-elements/references/conversation.md

@@ -198,9 +198,9 @@ import { ConversationDownload } from "@/components/ai-elements/conversation";
 
 | Prop | Type | Default | Description |
 |------|------|---------|-------------|
-| `messages` | `ConversationMessage[]` | Required | Array of messages to include in the download. |
+| `messages` | `UIMessage[]` | Required | Array of messages to include in the download. |
 | `filename` | `string` | - | The filename for the downloaded file. |
-| `formatMessage` | `(message: ConversationMessage, index: number) => string` | - | Custom function to format each message in the output. |
+| `formatMessage` | `(message: UIMessage, index: number) => string` | - | Custom function to format each message in the output. |
 | `...props` | `Omit<ComponentProps<typeof Button>, ` | - | Any other props are spread to the underlying shadcn/ui Button component. |
 
 ### `messagesToMarkdown`