refactor: replace message.content with message.parts (#868)

app/(chat)/api/chat/route.ts

@@ -1,5 +1,6 @@
 import {
-  type Message,
+  UIMessage,
+  appendResponseMessages,
   createDataStreamResponse,
   smoothStream,
   streamText,
@@ -15,7 +16,7 @@ import {
 import {
   generateUUID,
   getMostRecentUserMessage,
-  sanitizeResponseMessages,
+  getTrailingMessageId,
 } from '@/lib/utils';
 import { generateTitleFromUserMessage } from '../../actions';
 import { createDocument } from '@/lib/ai/tools/create-document';
@@ -23,7 +24,6 @@ import { updateDocument } from '@/lib/ai/tools/update-document';
 import { requestSuggestions } from '@/lib/ai/tools/request-suggestions';
 import { getWeather } from '@/lib/ai/tools/get-weather';
 import { isProductionEnvironment } from '@/lib/constants';
-import { NextResponse } from 'next/server';
 import { myProvider } from '@/lib/ai/providers';
 
 export const maxDuration = 60;
@@ -36,7 +36,7 @@ export async function POST(request: Request) {
       selectedChatModel,
     }: {
       id: string;
-      messages: Array<Message>;
+      messages: Array<UIMessage>;
       selectedChatModel: string;
     } = await request.json();
 
@@ -67,7 +67,16 @@ export async function POST(request: Request) {
     }
 
     await saveMessages({
-      messages: [{ ...userMessage, createdAt: new Date(), chatId: id }],
+      messages: [
+        {
+          chatId: id,
+          id: userMessage.id,
+          role: 'user',
+          parts: userMessage.parts,
+          attachments: userMessage.experimental_attachments ?? [],
+          createdAt: new Date(),
+        },
+      ],
     });
 
     return createDataStreamResponse({
@@ -97,24 +106,36 @@ export async function POST(request: Request) {
               dataStream,
             }),
           },
-          onFinish: async ({ response, reasoning }) => {
+          onFinish: async ({ response }) => {
             if (session.user?.id) {
               try {
-                const sanitizedResponseMessages = sanitizeResponseMessages({
+                const assistantId = getTrailingMessageId({
-                  messages: response.messages,
+                  messages: response.messages.filter(
-                  reasoning,
+                    (message) => message.role === 'assistant',
+                  ),
+                });
+
+                if (!assistantId) {
+                  throw new Error('No assistant message found!');
+                }
+
+                const [, assistantMessage] = appendResponseMessages({
+                  messages: [userMessage],
+                  responseMessages: response.messages,
                 });
 
                 await saveMessages({
-                  messages: sanitizedResponseMessages.map((message) => {
+                  messages: [
-                    return {
+                    {
-                      id: message.id,
+                      id: assistantId,
                       chatId: id,
-                      role: message.role,
+                      role: assistantMessage.role,
-                      content: message.content,
+                      parts: assistantMessage.parts,
+                      attachments:
+                        assistantMessage.experimental_attachments ?? [],
                       createdAt: new Date(),
-                    };
+                    },
-                  }),
+                  ],
                 });
               } catch (error) {
                 console.error('Failed to save chat');
@@ -138,7 +159,9 @@ export async function POST(request: Request) {
       },
     });
   } catch (error) {
-    return NextResponse.json({ error }, { status: 400 });
+    return new Response('An error occurred while processing your request!', {
+      status: 404,
+    });
   }
 }
 
@@ -167,7 +190,7 @@ export async function DELETE(request: Request) {
 
     return new Response('Chat deleted', { status: 200 });
   } catch (error) {
-    return new Response('An error occurred while processing your request', {
+    return new Response('An error occurred while processing your request!', {
       status: 500,
     });
   }

app/(chat)/chat/[id]/page.tsx

@@ -4,9 +4,10 @@ import { notFound } from 'next/navigation';
 import { auth } from '@/app/(auth)/auth';
 import { Chat } from '@/components/chat';
 import { getChatById, getMessagesByChatId } from '@/lib/db/queries';
-import { convertToUIMessages } from '@/lib/utils';
 import { DataStreamHandler } from '@/components/data-stream-handler';
 import { DEFAULT_CHAT_MODEL } from '@/lib/ai/models';
+import { DBMessage } from '@/lib/db/schema';
+import { Attachment, UIMessage } from 'ai';
 
 export default async function Page(props: { params: Promise<{ id: string }> }) {
   const params = await props.params;
@@ -33,6 +34,19 @@ export default async function Page(props: { params: Promise<{ id: string }> }) {
     id,
   });
 
+  function convertToUIMessages(messages: Array<DBMessage>): Array<UIMessage> {
+    return messages.map((message) => ({
+      id: message.id,
+      parts: message.parts as UIMessage['parts'],
+      role: message.role as UIMessage['role'],
+      // Note: content will soon be deprecated in @ai-sdk/react
+      content: '',
+      createdAt: message.createdAt,
+      experimental_attachments:
+        (message.attachments as Array<Attachment>) ?? [],
+    }));
+  }
+
   const cookieStore = await cookies();
   const chatModelFromCookie = cookieStore.get('chat-model');
 

app/layout.tsx

@@ -56,7 +56,7 @@ export default async function RootLayout({
           }}
         />
       </head>
-      <body className="antialiased">
+      <body className="">
         <ThemeProvider
           attribute="class"
           defaultTheme="system"

components/artifact-messages.tsx

@@ -1,7 +1,7 @@
 import { PreviewMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Vote } from '@/lib/db/schema';
-import { Message } from 'ai';
+import { UIMessage } from 'ai';
 import { memo } from 'react';
 import equal from 'fast-deep-equal';
 import { UIArtifact } from './artifact';
@@ -11,7 +11,7 @@ interface ArtifactMessagesProps {
   chatId: string;
   status: UseChatHelpers['status'];
   votes: Array<Vote> | undefined;
-  messages: Array<Message>;
+  messages: Array<UIMessage>;
   setMessages: UseChatHelpers['setMessages'];
   reload: UseChatHelpers['reload'];
   isReadonly: boolean;

components/artifact.tsx

@@ -1,4 +1,4 @@
-import type { Attachment, Message } from 'ai';
+import type { Attachment, UIMessage } from 'ai';
 import { formatDistance } from 'date-fns';
 import { AnimatePresence, motion } from 'framer-motion';
 import {
@@ -74,8 +74,8 @@ function PureArtifact({
   stop: UseChatHelpers['stop'];
   attachments: Array<Attachment>;
   setAttachments: Dispatch<SetStateAction<Array<Attachment>>>;
-  messages: Array<Message>;
+  messages: Array<UIMessage>;
-  setMessages: Dispatch<SetStateAction<Array<Message>>>;
+  setMessages: UseChatHelpers['setMessages'];
   votes: Array<Vote> | undefined;
   append: UseChatHelpers['append'];
   handleSubmit: UseChatHelpers['handleSubmit'];

components/chat.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import type { Attachment, Message } from 'ai';
+import type { Attachment, UIMessage } from 'ai';
 import { useChat } from '@ai-sdk/react';
 import { useState } from 'react';
 import useSWR, { useSWRConfig } from 'swr';
@@ -22,7 +22,7 @@ export function Chat({
   isReadonly,
 }: {
   id: string;
-  initialMessages: Array<Message>;
+  initialMessages: Array<UIMessage>;
   selectedChatModel: string;
   selectedVisibilityType: VisibilityType;
   isReadonly: boolean;

components/message-actions.tsx

@@ -32,8 +32,6 @@ export function PureMessageActions({
 
   if (isLoading) return null;
   if (message.role === 'user') return null;
-  if (message.toolInvocations && message.toolInvocations.length > 0)
-    return null;
 
   return (
     <TooltipProvider delayDuration={0}>
@@ -44,7 +42,18 @@ export function PureMessageActions({
               className="py-1 px-2 h-fit text-muted-foreground"
               variant="outline"
               onClick={async () => {
-                await copyToClipboard(message.content as string);
+                const textFromParts = message.parts
+                  ?.filter((part) => part.type === 'text')
+                  .map((part) => part.text)
+                  .join('\n')
+                  .trim();
+
+                if (!textFromParts) {
+                  toast.error("There's no text to copy!");
+                  return;
+                }
+
+                await copyToClipboard(textFromParts);
                 toast.success('Copied to clipboard!');
               }}
             >

components/message-editor.tsx

@@ -5,17 +5,13 @@ import { Button } from './ui/button';
 import { Dispatch, SetStateAction, useEffect, useRef, useState } from 'react';
 import { Textarea } from './ui/textarea';
 import { deleteTrailingMessages } from '@/app/(chat)/actions';
-import { toast } from 'sonner';
+import { UseChatHelpers } from '@ai-sdk/react';
 
 export type MessageEditorProps = {
   message: Message;
   setMode: Dispatch<SetStateAction<'view' | 'edit'>>;
-  setMessages: (
+  setMessages: UseChatHelpers['setMessages'];
-    messages: Message[] | ((messages: Message[]) => Message[]),
+  reload: UseChatHelpers['reload'];
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
 };
 
 export function MessageEditor({
@@ -79,6 +75,7 @@ export function MessageEditor({
               id: message.id,
             });
 
+            // @ts-expect-error todo: support UIMessage in setMessages
             setMessages((messages) => {
               const index = messages.findIndex((m) => m.id === message.id);
 
@@ -86,6 +83,7 @@ export function MessageEditor({
                 const updatedMessage = {
                   ...message,
                   content: draftContent,
+                  parts: [{ type: 'text', text: draftContent }],
                 };
 
                 return [...messages.slice(0, index), updatedMessage];

components/message.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import type { ChatRequestOptions, Message } from 'ai';
+import type { UIMessage } from 'ai';
 import cx from 'classnames';
 import { AnimatePresence, motion } from 'framer-motion';
 import { memo, useState } from 'react';
@@ -18,6 +18,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 import { MessageEditor } from './message-editor';
 import { DocumentPreview } from './document-preview';
 import { MessageReasoning } from './message-reasoning';
+import { UseChatHelpers } from '@ai-sdk/react';
 
 const PurePreviewMessage = ({
   chatId,
@@ -29,15 +30,11 @@ const PurePreviewMessage = ({
   isReadonly,
 }: {
   chatId: string;
-  message: Message;
+  message: UIMessage;
   vote: Vote | undefined;
   isLoading: boolean;
-  setMessages: (
+  setMessages: UseChatHelpers['setMessages'];
-    messages: Message[] | ((messages: Message[]) => Message[]),
+  reload: UseChatHelpers['reload'];
-  ) => void;
-  reload: (
-    chatRequestOptions?: ChatRequestOptions,
-  ) => Promise<string | null | undefined>;
   isReadonly: boolean;
 }) => {
   const [mode, setMode] = useState<'view' | 'edit'>('view');
@@ -83,96 +80,79 @@ const PurePreviewMessage = ({
               </div>
             )}
 
-            {message.reasoning && (
+            {message.parts?.map((part, index) => {
-              <MessageReasoning
+              const { type } = part;
-                isLoading={isLoading}
+              const key = `message-${message.id}-part-${index}`;
-                reasoning={message.reasoning}
-              />
-            )}
 
-            {(message.content || message.reasoning) && mode === 'view' && (
+              if (type === 'reasoning') {
-              <div
+                return (
-                data-testid="message-content"
+                  <MessageReasoning
-                className="flex flex-row gap-2 items-start"
+                    key={key}
-              >
+                    isLoading={isLoading}
-                {message.role === 'user' && !isReadonly && (
+                    reasoning={part.reasoning}
-                  <Tooltip>
+                  />
-                    <TooltipTrigger asChild>
+                );
-                      <Button
+              }
-                        data-testid={`message-edit`}
+
-                        variant="ghost"
+              if (type === 'text') {
-                        className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
+                if (mode === 'view') {
-                        onClick={() => {
+                  return (
-                          setMode('edit');
+                    <div key={key} className="flex flex-row gap-2 items-start">
-                        }}
+                      {message.role === 'user' && !isReadonly && (
+                        <Tooltip>
+                          <TooltipTrigger asChild>
+                            <Button
+                              data-testid="message-edit-button"
+                              variant="ghost"
+                              className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
+                              onClick={() => {
+                                setMode('edit');
+                              }}
+                            >
+                              <PencilEditIcon />
+                            </Button>
+                          </TooltipTrigger>
+                          <TooltipContent>Edit message</TooltipContent>
+                        </Tooltip>
+                      )}
+
+                      <div
+                        data-testid="message-content"
+                        className={cn('flex flex-col gap-4', {
+                          'bg-primary text-primary-foreground px-3 py-2 rounded-xl':
+                            message.role === 'user',
+                        })}
                       >
-                        <PencilEditIcon />
+                        <Markdown>{part.text}</Markdown>
-                      </Button>
-                    </TooltipTrigger>
-                    <TooltipContent>Edit message</TooltipContent>
-                  </Tooltip>
-                )}
-
-                <div
-                  className={cn('flex flex-col gap-4', {
-                    'bg-primary text-primary-foreground px-3 py-2 rounded-xl':
-                      message.role === 'user',
-                  })}
-                >
-                  <Markdown>{message.content as string}</Markdown>
-                </div>
-              </div>
-            )}
-
-            {message.content && mode === 'edit' && (
-              <div className="flex flex-row gap-2 items-start">
-                <div className="size-8" />
-
-                <MessageEditor
-                  key={message.id}
-                  message={message}
-                  setMode={setMode}
-                  setMessages={setMessages}
-                  reload={reload}
-                />
-              </div>
-            )}
-
-            {message.toolInvocations && message.toolInvocations.length > 0 && (
-              <div className="flex flex-col gap-4">
-                {message.toolInvocations.map((toolInvocation) => {
-                  const { toolName, toolCallId, state, args } = toolInvocation;
-
-                  if (state === 'result') {
-                    const { result } = toolInvocation;
-
-                    return (
-                      <div key={toolCallId}>
-                        {toolName === 'getWeather' ? (
-                          <Weather weatherAtLocation={result} />
-                        ) : toolName === 'createDocument' ? (
-                          <DocumentPreview
-                            isReadonly={isReadonly}
-                            result={result}
-                          />
-                        ) : toolName === 'updateDocument' ? (
-                          <DocumentToolResult
-                            type="update"
-                            result={result}
-                            isReadonly={isReadonly}
-                          />
-                        ) : toolName === 'requestSuggestions' ? (
-                          <DocumentToolResult
-                            type="request-suggestions"
-                            result={result}
-                            isReadonly={isReadonly}
-                          />
-                        ) : (
-                          <pre>{JSON.stringify(result, null, 2)}</pre>
-                        )}
                       </div>
-                    );
+                    </div>
-                  }
+                  );
+                }
+
+                if (mode === 'edit') {
+                  return (
+                    <div key={key} className="flex flex-row gap-2 items-start">
+                      <div className="size-8" />
+
+                      <MessageEditor
+                        key={message.id}
+                        message={message}
+                        setMode={setMode}
+                        setMessages={setMessages}
+                        reload={reload}
+                      />
+                    </div>
+                  );
+                }
+              }
+
+              if (type === 'tool-invocation') {
+                const { toolInvocation } = part;
+                const { toolName, toolCallId, state } = toolInvocation;
+
+                if (state === 'call') {
+                  const { args } = toolInvocation;
+
                   return (
                     <div
                       key={toolCallId}
@@ -199,9 +179,40 @@ const PurePreviewMessage = ({
                       ) : null}
                     </div>
                   );
-                })}
+                }
-              </div>
+
-            )}
+                if (state === 'result') {
+                  const { result } = toolInvocation;
+
+                  return (
+                    <div key={toolCallId}>
+                      {toolName === 'getWeather' ? (
+                        <Weather weatherAtLocation={result} />
+                      ) : toolName === 'createDocument' ? (
+                        <DocumentPreview
+                          isReadonly={isReadonly}
+                          result={result}
+                        />
+                      ) : toolName === 'updateDocument' ? (
+                        <DocumentToolResult
+                          type="update"
+                          result={result}
+                          isReadonly={isReadonly}
+                        />
+                      ) : toolName === 'requestSuggestions' ? (
+                        <DocumentToolResult
+                          type="request-suggestions"
+                          result={result}
+                          isReadonly={isReadonly}
+                        />
+                      ) : (
+                        <pre>{JSON.stringify(result, null, 2)}</pre>
+                      )}
+                    </div>
+                  );
+                }
+              }
+            })}
 
             {!isReadonly && (
               <MessageActions
@@ -223,16 +234,8 @@ export const PreviewMessage = memo(
   PurePreviewMessage,
   (prevProps, nextProps) => {
     if (prevProps.isLoading !== nextProps.isLoading) return false;
-    if (prevProps.message.reasoning !== nextProps.message.reasoning)
+    if (prevProps.message.id !== nextProps.message.id) return false;
-      return false;
+    if (!equal(prevProps.message.parts, nextProps.message.parts)) return false;
-    if (prevProps.message.content !== nextProps.message.content) return false;
-    if (
-      !equal(
-        prevProps.message.toolInvocations,
-        nextProps.message.toolInvocations,
-      )
-    )
-      return false;
     if (!equal(prevProps.vote, nextProps.vote)) return false;
 
     return true;
@@ -264,7 +267,7 @@ export const ThinkingMessage = () => {
 
         <div className="flex flex-col gap-2 w-full">
           <div className="flex flex-col gap-4 text-muted-foreground">
-            Thinking...
+            Hmm...
           </div>
         </div>
       </div>

components/messages.tsx

@@ -1,4 +1,4 @@
-import { Message } from 'ai';
+import { UIMessage } from 'ai';
 import { PreviewMessage, ThinkingMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Overview } from './overview';
@@ -11,7 +11,7 @@ interface MessagesProps {
   chatId: string;
   status: UseChatHelpers['status'];
   votes: Array<Vote> | undefined;
-  messages: Array<Message>;
+  messages: Array<UIMessage>;
   setMessages: UseChatHelpers['setMessages'];
   reload: UseChatHelpers['reload'];
   isReadonly: boolean;

components/multimodal-input.tsx

@@ -1,11 +1,6 @@
 'use client';
 
-import type {
+import type { Attachment, Message } from 'ai';
-  Attachment,
-  ChatRequestOptions,
-  CreateMessage,
-  Message,
-} from 'ai';
 import cx from 'classnames';
 import type React from 'react';
 import {
@@ -21,8 +16,6 @@ import {
 import { toast } from 'sonner';
 import { useLocalStorage, useWindowSize } from 'usehooks-ts';
 
-import { sanitizeUIMessages } from '@/lib/utils';
-
 import { ArrowUpIcon, PaperclipIcon, StopIcon } from './icons';
 import { PreviewAttachment } from './preview-attachment';
 import { Button } from './ui/button';
@@ -324,7 +317,7 @@ function PureStopButton({
       onClick={(event) => {
         event.preventDefault();
         stop();
-        setMessages((messages) => sanitizeUIMessages(messages));
+        setMessages((messages) => messages);
       }}
     >
       <StopIcon size={14} />

components/toolbar.tsx

@@ -25,18 +25,8 @@ import {
   TooltipProvider,
   TooltipTrigger,
 } from '@/components/ui/tooltip';
-import { sanitizeUIMessages } from '@/lib/utils';
 
-import {
+import { ArrowUpIcon, StopIcon, SummarizeIcon } from './icons';
-  ArrowUpIcon,
-  CodeIcon,
-  LogsIcon,
-  MessageIcon,
-  PenIcon,
-  SparklesIcon,
-  StopIcon,
-  SummarizeIcon,
-} from './icons';
 import { artifactDefinitions, ArtifactKind } from './artifact';
 import { ArtifactToolbarItem } from './create-artifact';
 import { UseChatHelpers } from '@ai-sdk/react';
@@ -442,7 +432,7 @@ const PureToolbar = ({
             className="p-3"
             onClick={() => {
               stop();
-              setMessages((messages) => sanitizeUIMessages(messages));
+              setMessages((messages) => messages);
             }}
           >
             <StopIcon />

docs/04-migrate-to-parts.md

@@ -0,0 +1,243 @@
+# Migrate to Message Parts
+
+The release of [`@ai-sdk/react@1.1.10`](https://github.com/vercel/ai/pull/4670) introduced a new property called `parts` to messages in the `useChat` hook. We recommend rendering the messages using the `parts` property instead of the `content` property. The parts property supports different message types, including text, tool invocation, and tool result, and allows for more flexible and complex chat and agent-like user interfaces.
+
+You can read the API reference for the `parts` property [here](https://sdk.vercel.ai/docs/reference/ai-sdk-ui/use-chat#messages.ui-message.parts).
+
+## Migrating your existing project to use the `parts` property
+
+Your existing project must already have messages stored in the database. To migrate your messages to use `parts`, you will have to create new tables for `Message` and `Vote`, backfill the new tables with transformed messages, and delete (optional) the old tables.
+
+These are the following steps:
+1. Create tables `Message_v2` and `Vote_v2` with updated schemas at `/lib/db/schema.ts`
+2. Update the `Message` component at `/components/message.tsx` to use parts and render content.
+3. Run migration script at `src/lib/db/helpers/01-migrate-to-parts.ts`
+
+### 1. Creating Tables with Updated Schemas
+
+You will mark the earlier tables as deprecated and create two new tables, `Message_v2` and `Vote_v2`. Table `Message_v2` contains the updated schema that includes the `parts` property. Table `Vote_v2` does not contain schema changes but will contain votes that point to the transformed messages in `Message_v2`.
+
+Before creating the new tables, you will need to update the variables of existing tables to indicate that they are deprecated.
+
+```tsx title="/lib/db/schema.ts"
+export const messageDeprecated = pgTable("Message", {
+  id: uuid("id").primaryKey().notNull().defaultRandom(),
+  chatId: uuid("chatId")
+    .notNull()
+    .references(() => chat.id),
+  role: varchar("role").notNull(),
+  content: json("content").notNull(),
+  createdAt: timestamp("createdAt").notNull(),
+});
+
+export type MessageDeprecated = InferSelectModel<typeof messageDeprecated>;
+
+export const voteDeprecated = pgTable(
+  "Vote",
+  {
+    chatId: uuid("chatId")
+      .notNull()
+      .references(() => chat.id),
+    messageId: uuid("messageId")
+      .notNull()
+      .references(() => messageDeprecated.id),
+    isUpvoted: boolean("isUpvoted").notNull(),
+  },
+  (table) => {
+    return {
+      pk: primaryKey({ columns: [table.chatId, table.messageId] }),
+    };
+  },
+);
+
+export type VoteDeprecated = InferSelectModel<typeof voteDeprecated>;
+```
+
+After deprecating the current table schemas, you can now proceed to create schemas for the new tables.
+
+```ts title="/lib/db/schema.ts"
+export const message = pgTable("Message_v2", {
+  id: uuid("id").primaryKey().notNull().defaultRandom(),
+  chatId: uuid("chatId")
+    .notNull()
+    .references(() => chat.id),
+  role: varchar("role").notNull(),
+  parts: json("parts").notNull(),
+  attachments: json("attachments").notNull(),
+  createdAt: timestamp("createdAt").notNull(),
+});
+
+export type Message = InferSelectModel<typeof message>;
+
+export const vote = pgTable(
+  "Vote_v2",
+  {
+    chatId: uuid("chatId")
+      .notNull()
+      .references(() => chat.id),
+    messageId: uuid("messageId")
+      .notNull()
+      .references(() => message.id),
+    isUpvoted: boolean("isUpvoted").notNull(),
+  },
+  (table) => {
+    return {
+      pk: primaryKey({ columns: [table.chatId, table.messageId] }),
+    };
+  },
+);
+
+export type Vote = InferSelectModel<typeof vote>;
+```
+
+### 2. Updating the Message Component
+
+Previously you were using content types to render messages and tool invocations. Now you will use the `parts` property to render messages and tool invocations.
+
+```tsx title="components/message.tsx"
+{message.parts?.map((part, index) => {
+  const { type } = part;
+  const key = `message-${message.id}-part-${index}`;
+
+  if (type === "reasoning") {
+    return (
+      <MessageReasoning
+        key={key}
+        isLoading={isLoading}
+        reasoning={part.reasoning}
+      />
+    );
+  }
+
+  if (type === "text") {
+    if (mode === "view") {
+      return (
+        <div key={key} className="flex flex-row gap-2 items-start">
+          {message.role === "user" && !isReadonly && (
+            <Tooltip>
+              <TooltipTrigger asChild>
+                <Button
+                  data-testid="message-edit-button"
+                  variant="ghost"
+                  className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
+                  onClick={() => {
+                    setMode("edit");
+                  }}
+                >
+                  <PencilEditIcon />
+                </Button>
+              </TooltipTrigger>
+              <TooltipContent>Edit message</TooltipContent>
+            </Tooltip>
+          )}
+
+          <div
+            data-testid="message-content"
+            className={cn("flex flex-col gap-4", {
+              "bg-primary text-primary-foreground px-3 py-2 rounded-xl":
+                message.role === "user",
+            })}
+          >
+            <Markdown>{part.text}</Markdown>
+          </div>
+        </div>
+      );
+    }
+
+    if (mode === "edit") {
+      return (
+        <div key={key} className="flex flex-row gap-2 items-start">
+          <div className="size-8" />
+
+          <MessageEditor
+            key={message.id}
+            message={message}
+            setMode={setMode}
+            setMessages={setMessages}
+            reload={reload}
+          />
+        </div>
+      );
+    }
+  }
+
+  if (type === "tool-invocation") {
+    const { toolInvocation } = part;
+    const { toolName, toolCallId, state } = toolInvocation;
+
+    if (state === "call") {
+      const { args } = toolInvocation;
+
+      return (
+        <div
+          key={toolCallId}
+          className={cx({
+            skeleton: ["getWeather"].includes(toolName),
+          })}
+        >
+          {toolName === "getWeather" ? (
+            <Weather />
+          ) : toolName === "createDocument" ? (
+            <DocumentPreview isReadonly={isReadonly} args={args} />
+          ) : toolName === "updateDocument" ? (
+            <DocumentToolCall
+              type="update"
+              args={args}
+              isReadonly={isReadonly}
+            />
+          ) : toolName === "requestSuggestions" ? (
+            <DocumentToolCall
+              type="request-suggestions"
+              args={args}
+              isReadonly={isReadonly}
+            />
+          ) : null}
+        </div>
+      );
+    }
+
+    if (state === "result") {
+      const { result } = toolInvocation;
+
+      return (
+        <div key={toolCallId}>
+          {toolName === "getWeather" ? (
+            <Weather weatherAtLocation={result} />
+          ) : toolName === "createDocument" ? (
+            <DocumentPreview
+              isReadonly={isReadonly}
+              result={result}
+            />
+          ) : toolName === "updateDocument" ? (
+            <DocumentToolResult
+              type="update"
+              result={result}
+              isReadonly={isReadonly}
+            />
+          ) : toolName === "requestSuggestions" ? (
+            <DocumentToolResult
+              type="request-suggestions"
+              result={result}
+              isReadonly={isReadonly}
+            />
+          ) : (
+            <pre>{JSON.stringify(result, null, 2)}</pre>
+          )}
+        </div>
+      );
+    }
+  }
+})}
+```
+
+### 3. Running the Migration Script
+
+At this point, you can deploy your application so new messages can be stored in the new format.
+
+To restore messages of previous chat conversations, you can run the following script that applies a transformation to the old messages and stores them in the new format.
+
+```zsh title="shell"
+pnpm exec tsx lib/db/helpers/01-core-to-parts.ts
+```
+
+This script will take some time to complete based on the number of messages to be migrated. After completion, you can verify that the messages have been successfully migrated by checking the database.

lib/db/helpers/01-core-to-parts.ts

@@ -0,0 +1,196 @@
+import { config } from 'dotenv';
+import postgres from 'postgres';
+import {
+  chat,
+  message,
+  messageDeprecated,
+  vote,
+  voteDeprecated,
+} from '../schema';
+import { drizzle } from 'drizzle-orm/postgres-js';
+import { inArray } from 'drizzle-orm';
+import { appendResponseMessages, UIMessage } from 'ai';
+
+config({
+  path: '.env.local',
+});
+
+if (!process.env.POSTGRES_URL) {
+  throw new Error('POSTGRES_URL environment variable is not set');
+}
+
+const client = postgres(process.env.POSTGRES_URL);
+const db = drizzle(client);
+
+const BATCH_SIZE = 50; // Process 10 chats at a time
+const INSERT_BATCH_SIZE = 100; // Insert 100 messages at a time
+
+type NewMessageInsert = {
+  id: string;
+  chatId: string;
+  parts: any[];
+  role: string;
+  attachments: any[];
+  createdAt: Date;
+};
+
+type NewVoteInsert = {
+  messageId: string;
+  chatId: string;
+  isUpvoted: boolean;
+};
+
+async function createNewTable() {
+  const chats = await db.select().from(chat);
+  let processedCount = 0;
+
+  // Process chats in batches
+  for (let i = 0; i < chats.length; i += BATCH_SIZE) {
+    const chatBatch = chats.slice(i, i + BATCH_SIZE);
+    const chatIds = chatBatch.map((chat) => chat.id);
+
+    // Fetch all messages and votes for the current batch of chats in bulk
+    const allMessages = await db
+      .select()
+      .from(messageDeprecated)
+      .where(inArray(messageDeprecated.chatId, chatIds));
+
+    const allVotes = await db
+      .select()
+      .from(voteDeprecated)
+      .where(inArray(voteDeprecated.chatId, chatIds));
+
+    // Prepare batches for insertion
+    const newMessagesToInsert: NewMessageInsert[] = [];
+    const newVotesToInsert: NewVoteInsert[] = [];
+
+    // Process each chat in the batch
+    for (const chat of chatBatch) {
+      processedCount++;
+      console.info(`Processed ${processedCount}/${chats.length} chats`);
+
+      // Filter messages and votes for this specific chat
+      const messages = allMessages.filter((msg) => msg.chatId === chat.id);
+      const votes = allVotes.filter((v) => v.chatId === chat.id);
+
+      // Group messages into sections
+      const messageSection: Array<UIMessage> = [];
+      const messageSections: Array<Array<UIMessage>> = [];
+
+      for (const message of messages) {
+        const { role } = message;
+
+        if (role === 'user' && messageSection.length > 0) {
+          messageSections.push([...messageSection]);
+          messageSection.length = 0;
+        }
+
+        // @ts-expect-error message.content has different type
+        messageSection.push(message);
+      }
+
+      if (messageSection.length > 0) {
+        messageSections.push([...messageSection]);
+      }
+
+      // Process each message section
+      for (const section of messageSections) {
+        const [userMessage, ...assistantMessages] = section;
+
+        const [firstAssistantMessage] = assistantMessages;
+
+        try {
+          const uiSection = appendResponseMessages({
+            messages: [userMessage],
+            // @ts-expect-error: message.content has different type
+            responseMessages: assistantMessages,
+            _internal: {
+              currentDate: () => firstAssistantMessage.createdAt ?? new Date(),
+            },
+          });
+
+          const projectedUISection = uiSection
+            .map((message) => {
+              if (message.role === 'user') {
+                return {
+                  id: message.id,
+                  chatId: chat.id,
+                  parts: [{ type: 'text', text: message.content }],
+                  role: message.role,
+                  createdAt: message.createdAt,
+                  attachments: [],
+                } as NewMessageInsert;
+              } else if (message.role === 'assistant') {
+                return {
+                  id: message.id,
+                  chatId: chat.id,
+                  parts: message.parts || [],
+                  role: message.role,
+                  createdAt: message.createdAt,
+                  attachments: [],
+                } as NewMessageInsert;
+              }
+              return null;
+            })
+            .filter((msg): msg is NewMessageInsert => msg !== null);
+
+          // Add messages to batch
+          for (const msg of projectedUISection) {
+            newMessagesToInsert.push(msg);
+
+            if (msg.role === 'assistant') {
+              const voteByMessage = votes.find((v) => v.messageId === msg.id);
+              if (voteByMessage) {
+                newVotesToInsert.push({
+                  messageId: msg.id,
+                  chatId: msg.chatId,
+                  isUpvoted: voteByMessage.isUpvoted,
+                });
+              }
+            }
+          }
+        } catch (error) {
+          console.error(`Error processing chat ${chat.id}: ${error}`);
+        }
+      }
+    }
+
+    // Batch insert messages
+    for (let j = 0; j < newMessagesToInsert.length; j += INSERT_BATCH_SIZE) {
+      const messageBatch = newMessagesToInsert.slice(j, j + INSERT_BATCH_SIZE);
+      if (messageBatch.length > 0) {
+        // Ensure all required fields are present
+        const validMessageBatch = messageBatch.map((msg) => ({
+          id: msg.id,
+          chatId: msg.chatId,
+          parts: msg.parts,
+          role: msg.role,
+          attachments: msg.attachments,
+          createdAt: msg.createdAt,
+        }));
+
+        await db.insert(message).values(validMessageBatch);
+      }
+    }
+
+    // Batch insert votes
+    for (let j = 0; j < newVotesToInsert.length; j += INSERT_BATCH_SIZE) {
+      const voteBatch = newVotesToInsert.slice(j, j + INSERT_BATCH_SIZE);
+      if (voteBatch.length > 0) {
+        await db.insert(vote).values(voteBatch);
+      }
+    }
+  }
+
+  console.info(`Migration completed: ${processedCount} chats processed`);
+}
+
+createNewTable()
+  .then(() => {
+    console.info('Script completed successfully');
+    process.exit(0);
+  })
+  .catch((error) => {
+    console.error('Script failed:', error);
+    process.exit(1);
+  });

lib/db/migrations/0005_wooden_whistler.sql

@@ -0,0 +1,33 @@
+CREATE TABLE IF NOT EXISTS "Message_v2" (
+	"id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
+	"chatId" uuid NOT NULL,
+	"role" varchar NOT NULL,
+	"parts" json NOT NULL,
+	"attachments" json NOT NULL,
+	"createdAt" timestamp NOT NULL
+);
+--> statement-breakpoint
+CREATE TABLE IF NOT EXISTS "Vote_v2" (
+	"chatId" uuid NOT NULL,
+	"messageId" uuid NOT NULL,
+	"isUpvoted" boolean NOT NULL,
+	CONSTRAINT "Vote_v2_chatId_messageId_pk" PRIMARY KEY("chatId","messageId")
+);
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "Message_v2" ADD CONSTRAINT "Message_v2_chatId_Chat_id_fk" FOREIGN KEY ("chatId") REFERENCES "public"."Chat"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "Vote_v2" ADD CONSTRAINT "Vote_v2_chatId_Chat_id_fk" FOREIGN KEY ("chatId") REFERENCES "public"."Chat"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;
+--> statement-breakpoint
+DO $$ BEGIN
+ ALTER TABLE "Vote_v2" ADD CONSTRAINT "Vote_v2_messageId_Message_v2_id_fk" FOREIGN KEY ("messageId") REFERENCES "public"."Message_v2"("id") ON DELETE no action ON UPDATE no action;
+EXCEPTION
+ WHEN duplicate_object THEN null;
+END $$;

lib/db/migrations/meta/0005_snapshot.json

@@ -0,0 +1,515 @@
+{
+  "id": "c6c102e6-b64e-4f0c-a7a6-32df758de437",
+  "prevId": "30ad8233-1432-428b-93fc-2bb1ba867ff1",
+  "version": "7",
+  "dialect": "postgresql",
+  "tables": {
+    "public.Chat": {
+      "name": "Chat",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "title": {
+          "name": "title",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "visibility": {
+          "name": "visibility",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "'private'"
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Chat_userId_User_id_fk": {
+          "name": "Chat_userId_User_id_fk",
+          "tableFrom": "Chat",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Document": {
+      "name": "Document",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "title": {
+          "name": "title",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "content": {
+          "name": "content",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "text": {
+          "name": "text",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "'text'"
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Document_userId_User_id_fk": {
+          "name": "Document_userId_User_id_fk",
+          "tableFrom": "Document",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Document_id_createdAt_pk": {
+          "name": "Document_id_createdAt_pk",
+          "columns": [
+            "id",
+            "createdAt"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.Message_v2": {
+      "name": "Message_v2",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "role": {
+          "name": "role",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "parts": {
+          "name": "parts",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "attachments": {
+          "name": "attachments",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Message_v2_chatId_Chat_id_fk": {
+          "name": "Message_v2_chatId_Chat_id_fk",
+          "tableFrom": "Message_v2",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Message": {
+      "name": "Message",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "role": {
+          "name": "role",
+          "type": "varchar",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "content": {
+          "name": "content",
+          "type": "json",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Message_chatId_Chat_id_fk": {
+          "name": "Message_chatId_Chat_id_fk",
+          "tableFrom": "Message",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Suggestion": {
+      "name": "Suggestion",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "documentId": {
+          "name": "documentId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "documentCreatedAt": {
+          "name": "documentCreatedAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "originalText": {
+          "name": "originalText",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "suggestedText": {
+          "name": "suggestedText",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "description": {
+          "name": "description",
+          "type": "text",
+          "primaryKey": false,
+          "notNull": false
+        },
+        "isResolved": {
+          "name": "isResolved",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true,
+          "default": false
+        },
+        "userId": {
+          "name": "userId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "createdAt": {
+          "name": "createdAt",
+          "type": "timestamp",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Suggestion_userId_User_id_fk": {
+          "name": "Suggestion_userId_User_id_fk",
+          "tableFrom": "Suggestion",
+          "tableTo": "User",
+          "columnsFrom": [
+            "userId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk": {
+          "name": "Suggestion_documentId_documentCreatedAt_Document_id_createdAt_fk",
+          "tableFrom": "Suggestion",
+          "tableTo": "Document",
+          "columnsFrom": [
+            "documentId",
+            "documentCreatedAt"
+          ],
+          "columnsTo": [
+            "id",
+            "createdAt"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Suggestion_id_pk": {
+          "name": "Suggestion_id_pk",
+          "columns": [
+            "id"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.User": {
+      "name": "User",
+      "schema": "",
+      "columns": {
+        "id": {
+          "name": "id",
+          "type": "uuid",
+          "primaryKey": true,
+          "notNull": true,
+          "default": "gen_random_uuid()"
+        },
+        "email": {
+          "name": "email",
+          "type": "varchar(64)",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "password": {
+          "name": "password",
+          "type": "varchar(64)",
+          "primaryKey": false,
+          "notNull": false
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {},
+      "compositePrimaryKeys": {},
+      "uniqueConstraints": {}
+    },
+    "public.Vote_v2": {
+      "name": "Vote_v2",
+      "schema": "",
+      "columns": {
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "messageId": {
+          "name": "messageId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "isUpvoted": {
+          "name": "isUpvoted",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Vote_v2_chatId_Chat_id_fk": {
+          "name": "Vote_v2_chatId_Chat_id_fk",
+          "tableFrom": "Vote_v2",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Vote_v2_messageId_Message_v2_id_fk": {
+          "name": "Vote_v2_messageId_Message_v2_id_fk",
+          "tableFrom": "Vote_v2",
+          "tableTo": "Message_v2",
+          "columnsFrom": [
+            "messageId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Vote_v2_chatId_messageId_pk": {
+          "name": "Vote_v2_chatId_messageId_pk",
+          "columns": [
+            "chatId",
+            "messageId"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    },
+    "public.Vote": {
+      "name": "Vote",
+      "schema": "",
+      "columns": {
+        "chatId": {
+          "name": "chatId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "messageId": {
+          "name": "messageId",
+          "type": "uuid",
+          "primaryKey": false,
+          "notNull": true
+        },
+        "isUpvoted": {
+          "name": "isUpvoted",
+          "type": "boolean",
+          "primaryKey": false,
+          "notNull": true
+        }
+      },
+      "indexes": {},
+      "foreignKeys": {
+        "Vote_chatId_Chat_id_fk": {
+          "name": "Vote_chatId_Chat_id_fk",
+          "tableFrom": "Vote",
+          "tableTo": "Chat",
+          "columnsFrom": [
+            "chatId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        },
+        "Vote_messageId_Message_id_fk": {
+          "name": "Vote_messageId_Message_id_fk",
+          "tableFrom": "Vote",
+          "tableTo": "Message",
+          "columnsFrom": [
+            "messageId"
+          ],
+          "columnsTo": [
+            "id"
+          ],
+          "onDelete": "no action",
+          "onUpdate": "no action"
+        }
+      },
+      "compositePrimaryKeys": {
+        "Vote_chatId_messageId_pk": {
+          "name": "Vote_chatId_messageId_pk",
+          "columns": [
+            "chatId",
+            "messageId"
+          ]
+        }
+      },
+      "uniqueConstraints": {}
+    }
+  },
+  "enums": {},
+  "schemas": {},
+  "sequences": {},
+  "_meta": {
+    "columns": {},
+    "schemas": {},
+    "tables": {}
+  }
+}

lib/db/migrations/meta/_journal.json

@@ -36,6 +36,13 @@
       "when": 1733945232355,
       "tag": "0004_odd_slayback",
       "breakpoints": true
+    },
+    {
+      "idx": 5,
+      "version": "7",
+      "when": 1741934630596,
+      "tag": "0005_wooden_whistler",
+      "breakpoints": true
     }
   ]
 }

lib/db/queries.ts

@@ -12,9 +12,9 @@ import {
   document,
   type Suggestion,
   suggestion,
-  type Message,
   message,
   vote,
+  type DBMessage,
 } from './schema';
 import { ArtifactKind } from '@/components/artifact';
 
@@ -104,7 +104,11 @@ export async function getChatById({ id }: { id: string }) {
   }
 }
 
-export async function saveMessages({ messages }: { messages: Array<Message> }) {
+export async function saveMessages({
+  messages,
+}: {
+  messages: Array<DBMessage>;
+}) {
   try {
     return await db.insert(message).values(messages);
   } catch (error) {

lib/db/schema.ts

@@ -33,7 +33,9 @@ export const chat = pgTable('Chat', {
 
 export type Chat = InferSelectModel<typeof chat>;
 
-export const message = pgTable('Message', {
+// DEPRECATED: The following schema is deprecated and will be removed in the future.
+// Read the migration guide at https://github.com/vercel/ai-chatbot/blob/main/docs/04-migrate-to-parts.md
+export const messageDeprecated = pgTable('Message', {
   id: uuid('id').primaryKey().notNull().defaultRandom(),
   chatId: uuid('chatId')
     .notNull()
@@ -43,10 +45,45 @@ export const message = pgTable('Message', {
   createdAt: timestamp('createdAt').notNull(),
 });
 
-export type Message = InferSelectModel<typeof message>;
+export type MessageDeprecated = InferSelectModel<typeof messageDeprecated>;
+
+export const message = pgTable('Message_v2', {
+  id: uuid('id').primaryKey().notNull().defaultRandom(),
+  chatId: uuid('chatId')
+    .notNull()
+    .references(() => chat.id),
+  role: varchar('role').notNull(),
+  parts: json('parts').notNull(),
+  attachments: json('attachments').notNull(),
+  createdAt: timestamp('createdAt').notNull(),
+});
+
+export type DBMessage = InferSelectModel<typeof message>;
+
+// DEPRECATED: The following schema is deprecated and will be removed in the future.
+// Read the migration guide at https://github.com/vercel/ai-chatbot/blob/main/docs/04-migrate-to-parts.md
+export const voteDeprecated = pgTable(
+  'Vote',
+  {
+    chatId: uuid('chatId')
+      .notNull()
+      .references(() => chat.id),
+    messageId: uuid('messageId')
+      .notNull()
+      .references(() => messageDeprecated.id),
+    isUpvoted: boolean('isUpvoted').notNull(),
+  },
+  (table) => {
+    return {
+      pk: primaryKey({ columns: [table.chatId, table.messageId] }),
+    };
+  },
+);
+
+export type VoteDeprecated = InferSelectModel<typeof voteDeprecated>;
 
 export const vote = pgTable(
-  'Vote',
+  'Vote_v2',
   {
     chatId: uuid('chatId')
       .notNull()

lib/utils.ts

@@ -5,11 +5,12 @@ import type {
   TextStreamPart,
   ToolInvocation,
   ToolSet,
+  UIMessage,
 } from 'ai';
 import { type ClassValue, clsx } from 'clsx';
 import { twMerge } from 'tailwind-merge';
 
-import type { Message as DBMessage, Document } from '@/lib/db/schema';
+import type { DBMessage, Document } from '@/lib/db/schema';
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs));
@@ -85,52 +86,6 @@ function addToolMessageToChat({
   });
 }
 
-export function convertToUIMessages(
-  messages: Array<DBMessage>,
-): Array<Message> {
-  return messages.reduce((chatMessages: Array<Message>, message) => {
-    if (message.role === 'tool') {
-      return addToolMessageToChat({
-        toolMessage: message as CoreToolMessage,
-        messages: chatMessages,
-      });
-    }
-
-    let textContent = '';
-    let reasoning: string | undefined = undefined;
-    const toolInvocations: Array<ToolInvocation> = [];
-
-    if (typeof message.content === 'string') {
-      textContent = message.content;
-    } else if (Array.isArray(message.content)) {
-      for (const content of message.content) {
-        if (content.type === 'text') {
-          textContent += content.text;
-        } else if (content.type === 'tool-call') {
-          toolInvocations.push({
-            state: 'call',
-            toolCallId: content.toolCallId,
-            toolName: content.toolName,
-            args: content.args,
-          });
-        } else if (content.type === 'reasoning') {
-          reasoning = content.reasoning;
-        }
-      }
-    }
-
-    chatMessages.push({
-      id: message.id,
-      role: message.role as Message['role'],
-      content: textContent,
-      reasoning,
-      toolInvocations,
-    });
-
-    return chatMessages;
-  }, []);
-}
-
 type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage;
 type ResponseMessage = ResponseMessageWithoutId & { id: string };
 
@@ -182,40 +137,7 @@ export function sanitizeResponseMessages({
   );
 }
 
-export function sanitizeUIMessages(messages: Array<Message>): Array<Message> {
+export function getMostRecentUserMessage(messages: Array<UIMessage>) {
-  const messagesBySanitizedToolInvocations = messages.map((message) => {
-    if (message.role !== 'assistant') return message;
-
-    if (!message.toolInvocations) return message;
-
-    const toolResultIds: Array<string> = [];
-
-    for (const toolInvocation of message.toolInvocations) {
-      if (toolInvocation.state === 'result') {
-        toolResultIds.push(toolInvocation.toolCallId);
-      }
-    }
-
-    const sanitizedToolInvocations = message.toolInvocations.filter(
-      (toolInvocation) =>
-        toolInvocation.state === 'result' ||
-        toolResultIds.includes(toolInvocation.toolCallId),
-    );
-
-    return {
-      ...message,
-      toolInvocations: sanitizedToolInvocations,
-    };
-  });
-
-  return messagesBySanitizedToolInvocations.filter(
-    (message) =>
-      message.content.length > 0 ||
-      (message.toolInvocations && message.toolInvocations.length > 0),
-  );
-}
-
-export function getMostRecentUserMessage(messages: Array<Message>) {
   const userMessages = messages.filter((message) => message.role === 'user');
   return userMessages.at(-1);
 }
@@ -229,3 +151,15 @@ export function getDocumentTimestampByIndex(
 
   return documents[index].createdAt;
 }
+
+export function getTrailingMessageId({
+  messages,
+}: {
+  messages: Array<ResponseMessage>;
+}): string | null {
+  const trailingMessage = messages.at(-1);
+
+  if (!trailingMessage) return null;
+
+  return trailingMessage.id;
+}

package.json

@@ -19,9 +19,9 @@
     "test": "export PLAYWRIGHT=True && pnpm exec playwright test --workers=4"
   },
   "dependencies": {
-    "@ai-sdk/fireworks": "0.1.12",
+    "@ai-sdk/fireworks": "0.1.16",
-    "@ai-sdk/openai": "1.2.0",
+    "@ai-sdk/openai": "1.2.5",
-    "@ai-sdk/react": "^1.1.20",
+    "@ai-sdk/react": "^1.1.23",
     "@codemirror/lang-javascript": "^6.2.2",
     "@codemirror/lang-python": "^6.1.6",
     "@codemirror/state": "^6.5.0",
@@ -40,7 +40,7 @@
     "@vercel/analytics": "^1.3.1",
     "@vercel/blob": "^0.24.1",
     "@vercel/postgres": "^0.10.0",
-    "ai": "4.1.50",
+    "ai": "4.1.61",
     "bcrypt-ts": "^5.0.2",
     "class-variance-authority": "^0.7.0",
     "classnames": "^2.5.1",

playwright.config.ts

@@ -44,9 +44,9 @@ export default defineConfig({
   },
 
   /* Configure global timeout for each test */
-  timeout: 30000,
+  timeout: 60 * 1000, // 30 seconds
   expect: {
-    timeout: 30000,
+    timeout: 60 * 1000,
   },
 
   /* Configure projects */

pnpm-lock.yaml

@@ -9,14 +9,14 @@ importers:
   .:
     dependencies:
       '@ai-sdk/fireworks':
-        specifier: 0.1.12
+        specifier: 0.1.16
-        version: 0.1.12(zod@3.23.8)
+        version: 0.1.16(zod@3.23.8)
       '@ai-sdk/openai':
-        specifier: 1.2.0
+        specifier: 1.2.5
-        version: 1.2.0(zod@3.23.8)
+        version: 1.2.5(zod@3.23.8)
       '@ai-sdk/react':
-        specifier: ^1.1.20
+        specifier: ^1.1.23
-        version: 1.1.20(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
+        version: 1.1.23(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
       '@codemirror/lang-javascript':
         specifier: ^6.2.2
         version: 6.2.2
@@ -72,8 +72,8 @@ importers:
         specifier: ^0.10.0
         version: 0.10.0
       ai:
-        specifier: 4.1.50
+        specifier: 4.1.61
-        version: 4.1.50(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
+        version: 4.1.61(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
       bcrypt-ts:
         specifier: ^5.0.2
         version: 5.0.2
@@ -258,26 +258,26 @@ importers:
 
 packages:
 
-  '@ai-sdk/fireworks@0.1.12':
+  '@ai-sdk/fireworks@0.1.16':
-    resolution: {integrity: sha512-13HGPMRDUvdS7d53stNiVMyQa2S0PiefymYjIaXwz8KhA3ZEfxTyzw+42OL84FdEuPGFMadKkXOdvWSW0YVUDg==}
+    resolution: {integrity: sha512-CQ0A/bVYUQ2VBLxH6RHOW+m+1LlJ9QNf/7hjYFWAmBR6nkwHbrKarNOX9d3Ba08HTotdqxIChORuU2km0LIxFw==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
 
-  '@ai-sdk/openai-compatible@0.1.12':
+  '@ai-sdk/openai-compatible@0.1.15':
-    resolution: {integrity: sha512-2bMhAEeiRz4lbW5ixjGjbPhwyqjtujkjLVpqqtqWvvUDvtUM3cw1go9pqWFgaNKSBDaXRUfi8mkAVrn1yRuY2A==}
+    resolution: {integrity: sha512-iuylARLylTkaVYXKdgA1GeGU2TwANaJ3RlhJqmN0cVO5uEt3qVKvedrzR3jrQIyy9hu3gngsiXfFBMW2UQ7Ntg==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
 
-  '@ai-sdk/openai@1.2.0':
+  '@ai-sdk/openai@1.2.5':
-    resolution: {integrity: sha512-tzxH6OxKL5ffts4zJPdziQSJGGpSrQcJmuSrE92jCt7pJ4PAU5Dx4tjNNFIU8lSfwarLnywejZEt3Fz0uQZZOQ==}
+    resolution: {integrity: sha512-COK7LzspgQQh5Yq070xfDdVMvp8WX592rXRaMaYNNqu1xpzahxDcM24aF9xgKYWuYH0UMoOw4UmWGwGxr6ygIg==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
 
-  '@ai-sdk/provider-utils@2.1.10':
+  '@ai-sdk/provider-utils@2.1.13':
-    resolution: {integrity: sha512-4GZ8GHjOFxePFzkl3q42AU0DQOtTQ5w09vmaWUf/pKFXJPizlnzKSUkF0f+VkapIUfDugyMqPMT1ge8XQzVI7Q==}
+    resolution: {integrity: sha512-kLjqsfOdONr6DGcGEntFYM1niXz1H05vyZNf9OAzK+KKKc64izyP4/q/9HX7W4+6g8hm6BnmKxu8vkr6FSOqDg==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
@@ -285,12 +285,12 @@ packages:
       zod:
         optional: true
 
-  '@ai-sdk/provider@1.0.9':
+  '@ai-sdk/provider@1.0.11':
-    resolution: {integrity: sha512-jie6ZJT2ZR0uVOVCDc9R2xCX5I/Dum/wEK28lx21PJx6ZnFAN9EzD2WsPhcDWfCgGx3OAZZ0GyM3CEobXpa9LA==}
+    resolution: {integrity: sha512-CPyImHGiT3svyfmvPvAFTianZzWFtm0qK82XjwlQIA1C3IQ2iku/PMQXi7aFyrX0TyMh3VTkJPB03tjU2VXVrw==}
     engines: {node: '>=18'}
 
-  '@ai-sdk/react@1.1.20':
+  '@ai-sdk/react@1.1.23':
-    resolution: {integrity: sha512-4QOM9fR9SryaRraybckDjrhl1O6XejqELdKmrM5g9y9eLnWAfjwF+W1aN0knkSHzbbjMqN77sy9B9yL8EuJbDw==}
+    resolution: {integrity: sha512-R+PG9ya0GLs6orzt+1MxmjrWFuZM0gVs+l8ihBr1u+42wwkVeojY4CAtQjW4nrfGTVbdJYkl5y+r/VKfjr42aQ==}
     engines: {node: '>=18'}
     peerDependencies:
       react: ^18 || ^19 || ^19.0.0-rc
@@ -301,8 +301,8 @@ packages:
       zod:
         optional: true
 
-  '@ai-sdk/ui-utils@1.1.16':
+  '@ai-sdk/ui-utils@1.1.19':
-    resolution: {integrity: sha512-jfblR2yZVISmNK2zyNzJZFtkgX57WDAUQXcmn3XUBJyo8LFsADu+/vYMn5AOyBi9qJT0RBk11PEtIxIqvByw3Q==}
+    resolution: {integrity: sha512-rDHy2uxlPMt3jjS9L6mBrsfhEInZ5BVoWevmD13fsAt2s/XWy2OwwKmgmUQkdLlY4mn/eyeYAfDGK8+5CbOAgg==}
     engines: {node: '>=18'}
     peerDependencies:
       zod: ^3.0.0
@@ -1651,8 +1651,8 @@ packages:
     engines: {node: '>=0.4.0'}
     hasBin: true
 
-  ai@4.1.50:
+  ai@4.1.61:
-    resolution: {integrity: sha512-YBNeemrJKDrxoBQd3V9aaxhKm5q5YyRcF7PZE7W0NmLuvsdva/1aQNYTAsxs47gQFdvqfYmlFy4B0E+356OlPA==}
+    resolution: {integrity: sha512-Y9SAyGJEeW23F6C7PSHZXYNEvbH2cqJm0rVW2AoeFaXFT13ttx8rAqs8wz2w466C1UB329yl5PXayFcHqofSEA==}
     engines: {node: '>=18'}
     peerDependencies:
       react: ^18 || ^19 || ^19.0.0-rc
@@ -3800,52 +3800,52 @@ packages:
 
 snapshots:
 
-  '@ai-sdk/fireworks@0.1.12(zod@3.23.8)':
+  '@ai-sdk/fireworks@0.1.16(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/openai-compatible': 0.1.12(zod@3.23.8)
+      '@ai-sdk/openai-compatible': 0.1.15(zod@3.23.8)
-      '@ai-sdk/provider': 1.0.9
+      '@ai-sdk/provider': 1.0.11
-      '@ai-sdk/provider-utils': 2.1.10(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.13(zod@3.23.8)
       zod: 3.23.8
 
-  '@ai-sdk/openai-compatible@0.1.12(zod@3.23.8)':
+  '@ai-sdk/openai-compatible@0.1.15(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.9
+      '@ai-sdk/provider': 1.0.11
-      '@ai-sdk/provider-utils': 2.1.10(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.13(zod@3.23.8)
       zod: 3.23.8
 
-  '@ai-sdk/openai@1.2.0(zod@3.23.8)':
+  '@ai-sdk/openai@1.2.5(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.9
+      '@ai-sdk/provider': 1.0.11
-      '@ai-sdk/provider-utils': 2.1.10(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.13(zod@3.23.8)
       zod: 3.23.8
 
-  '@ai-sdk/provider-utils@2.1.10(zod@3.23.8)':
+  '@ai-sdk/provider-utils@2.1.13(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.9
+      '@ai-sdk/provider': 1.0.11
       eventsource-parser: 3.0.0
       nanoid: 3.3.8
       secure-json-parse: 2.7.0
     optionalDependencies:
       zod: 3.23.8
 
-  '@ai-sdk/provider@1.0.9':
+  '@ai-sdk/provider@1.0.11':
     dependencies:
       json-schema: 0.4.0
 
-  '@ai-sdk/react@1.1.20(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)':
+  '@ai-sdk/react@1.1.23(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider-utils': 2.1.10(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.13(zod@3.23.8)
-      '@ai-sdk/ui-utils': 1.1.16(zod@3.23.8)
+      '@ai-sdk/ui-utils': 1.1.19(zod@3.23.8)
       swr: 2.2.5(react@19.0.0-rc-45804af1-20241021)
       throttleit: 2.1.0
     optionalDependencies:
       react: 19.0.0-rc-45804af1-20241021
       zod: 3.23.8
 
-  '@ai-sdk/ui-utils@1.1.16(zod@3.23.8)':
+  '@ai-sdk/ui-utils@1.1.19(zod@3.23.8)':
     dependencies:
-      '@ai-sdk/provider': 1.0.9
+      '@ai-sdk/provider': 1.0.11
-      '@ai-sdk/provider-utils': 2.1.10(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.13(zod@3.23.8)
       zod-to-json-schema: 3.24.1(zod@3.23.8)
     optionalDependencies:
       zod: 3.23.8
@@ -4953,13 +4953,14 @@ snapshots:
 
   acorn@8.14.0: {}
 
-  ai@4.1.50(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8):
+  ai@4.1.61(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8):
     dependencies:
-      '@ai-sdk/provider': 1.0.9
+      '@ai-sdk/provider': 1.0.11
-      '@ai-sdk/provider-utils': 2.1.10(zod@3.23.8)
+      '@ai-sdk/provider-utils': 2.1.13(zod@3.23.8)
-      '@ai-sdk/react': 1.1.20(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
+      '@ai-sdk/react': 1.1.23(react@19.0.0-rc-45804af1-20241021)(zod@3.23.8)
-      '@ai-sdk/ui-utils': 1.1.16(zod@3.23.8)
+      '@ai-sdk/ui-utils': 1.1.19(zod@3.23.8)
       '@opentelemetry/api': 1.9.0
+      eventsource-parser: 3.0.0
       jsondiffpatch: 0.6.0
     optionalDependencies:
       react: 19.0.0-rc-45804af1-20241021

tests/pages/artifact.ts

@@ -92,7 +92,7 @@ export class ArtifactPage {
       content,
       attachments,
       async edit(newMessage: string) {
-        await page.getByTestId('message-edit').click();
+        await page.getByTestId('message-edit-button').click();
         await page.getByTestId('message-editor').fill(newMessage);
         await page.getByTestId('message-editor-send-button').click();
         await expect(

tests/pages/chat.ts

@@ -170,7 +170,7 @@ export class ChatPage {
       content,
       attachments,
       async edit(newMessage: string) {
-        await page.getByTestId('message-edit').click();
+        await page.getByTestId('message-edit-button').click();
         await page.getByTestId('message-editor').fill(newMessage);
         await page.getByTestId('message-editor-send-button').click();
         await expect(