From 9f0dd1f1527b443a8df4cfcefd3f2795c877fd0a Mon Sep 17 00:00:00 2001
From: Nico Albanese <nico.albanese@vercel.com>
Date: Wed, 19 Feb 2025 12:26:41 -0500
Subject: [PATCH] refactor to message parts

---
 components/message.tsx | 161 ++++++++++++++++++++++-------------------
 1 file changed, 85 insertions(+), 76 deletions(-)

diff --git a/components/message.tsx b/components/message.tsx
index 1ebe1b933..10916e01c 100644
--- a/components/message.tsx
+++ b/components/message.tsx
@@ -3,14 +3,12 @@
 import type { ChatRequestOptions, Message } from 'ai';
 import cx from 'classnames';
 import { AnimatePresence, motion } from 'framer-motion';
-import { memo, useMemo, useState } from 'react';
+import { memo, useState } from 'react';
 
 import type { Vote } from '@/lib/db/schema';
 
 import { DocumentToolCall, DocumentToolResult } from './document';
 import {
-  ChevronDownIcon,
-  LoaderIcon,
   PencilEditIcon,
   SparklesIcon,
 } from './icons';
@@ -90,89 +88,99 @@ const PurePreviewMessage = ({
               </div>
             )}
 
-            {message.reasoning && (
-              <MessageReasoning
-                isLoading={isLoading}
-                reasoning={message.reasoning}
-              />
-            )}
-
-            {(message.content || message.reasoning) && mode === 'view' && (
-              <div className="flex flex-row gap-2 items-start">
-                {message.role === 'user' && !isReadonly && (
-                  <Tooltip>
-                    <TooltipTrigger asChild>
-                      <Button
-                        variant="ghost"
-                        className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
-                        onClick={() => {
-                          if (!user) {
-                            toast.error(
-                              'You must be signed in to edit messages!',
-                            );
-
-                            return;
-                          }
-
-                          setMode('edit');
-                        }}
+            {/* use message parts over content and toolInvocations directly */}
+            {message.parts?.map((p, i) => {
+              const key = `message-${message.id}-part-${i}`;
+              switch (p.type) {
+                case "text":
+                  if (mode === "view") {
+                    return (
+                      <div
+                        key={key}
+                        className="flex flex-row gap-2 items-start"
                       >
-                        <PencilEditIcon />
-                      </Button>
-                    </TooltipTrigger>
-                    <TooltipContent>Edit message</TooltipContent>
-                  </Tooltip>
-                )}
+                        {message.role === "user" && !isReadonly && (
+                          <Tooltip>
+                            <TooltipTrigger asChild>
+                              <Button
+                                variant="ghost"
+                                className="px-2 h-fit rounded-full text-muted-foreground opacity-0 group-hover/message:opacity-100"
+                                onClick={() => {
+                                  if (!user) {
+                                    toast.error(
+                                      "You must be signed in to edit messages!",
+                                    );
 
-                <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" />
+                                    return;
+                                  }
 
-                <MessageEditor
-                  key={message.id}
-                  message={message}
-                  setMode={setMode}
-                  setMessages={setMessages}
-                  reload={reload}
-                />
-              </div>
-            )}
+                                  setMode("edit");
+                                }}
+                              >
+                                <PencilEditIcon />
+                              </Button>
+                            </TooltipTrigger>
+                            <TooltipContent>Edit message</TooltipContent>
+                          </Tooltip>
+                        )}
 
-            {message.toolInvocations && message.toolInvocations.length > 0 && (
-              <div className="flex flex-col gap-4">
-                {message.toolInvocations.map((toolInvocation) => {
-                  const { toolName, toolCallId, state, args } = toolInvocation;
+                        <div
+                          className={cn("flex flex-col gap-4", {
+                            "bg-primary text-primary-foreground px-3 py-2 rounded-xl":
+                              message.role === "user",
+                          })}
+                        >
+                          <Markdown>{p.text}</Markdown>
+                        </div>
+                      </div>
+                    );
+                  }
+                  if (mode === "edit") {
+                    return (
+                      <div key={key} className="flex flex-row gap-2 items-start">
+                        <div className="size-8" />
 
-                  if (state === 'result') {
-                    const { result } = toolInvocation;
+                        <MessageEditor
+                          key={message.id}
+                          message={message}
+                          setMode={setMode}
+                          setMessages={setMessages}
+                          reload={reload}
+                        />
+                      </div>
+                    );
+                  }
+                  return null;
+                case "reasoning":
+                  return (
+                    <MessageReasoning
+                      key={key}
+                      isLoading={isLoading}
+                      reasoning={p.reasoning}
+                    />
+                  );
+                case "tool-invocation":
+                  const { toolName, toolCallId, state, args } =
+                    p.toolInvocation;
+                  if (state === "result") {
+                    const { result } = p.toolInvocation;
 
                     return (
                       <div key={toolCallId}>
-                        {toolName === 'getWeather' ? (
+                        {toolName === "getWeather" ? (
                           <Weather weatherAtLocation={result} />
-                        ) : toolName === 'createDocument' ? (
+                        ) : toolName === "createDocument" ? (
                           <DocumentPreview
                             isReadonly={isReadonly}
                             result={result}
                           />
-                        ) : toolName === 'updateDocument' ? (
+                        ) : toolName === "updateDocument" ? (
                           <DocumentToolResult
                             type="update"
                             result={result}
                             isReadonly={isReadonly}
                           />
-                        ) : toolName === 'requestSuggestions' ? (
+                        ) : toolName === "requestSuggestions" ? (
                           <DocumentToolResult
                             type="request-suggestions"
                             result={result}
@@ -188,20 +196,20 @@ const PurePreviewMessage = ({
                     <div
                       key={toolCallId}
                       className={cx({
-                        skeleton: ['getWeather'].includes(toolName),
+                        skeleton: ["getWeather"].includes(toolName),
                       })}
                     >
-                      {toolName === 'getWeather' ? (
+                      {toolName === "getWeather" ? (
                         <Weather />
-                      ) : toolName === 'createDocument' ? (
+                      ) : toolName === "createDocument" ? (
                         <DocumentPreview isReadonly={isReadonly} args={args} />
-                      ) : toolName === 'updateDocument' ? (
+                      ) : toolName === "updateDocument" ? (
                         <DocumentToolCall
                           type="update"
                           args={args}
                           isReadonly={isReadonly}
                         />
-                      ) : toolName === 'requestSuggestions' ? (
+                      ) : toolName === "requestSuggestions" ? (
                         <DocumentToolCall
                           type="request-suggestions"
                           args={args}
@@ -210,9 +218,10 @@ const PurePreviewMessage = ({
                       ) : null}
                     </div>
                   );
-                })}
-              </div>
-            )}
+                default:
+                  return null;
+              }
+            })}
 
             {!isReadonly && (
               <MessageActions