Merge branch 'vercel:main' into main

athrael-soju · Dec 5, 2024 · ce90ccb · ce90ccb
2 parents bbd1320 + 64d1aad
commit ce90ccb
Show file tree

Hide file tree

Showing 12 changed files with 302 additions and 25 deletions.
diff --git a/app/(chat)/actions.ts b/app/(chat)/actions.ts
@@ -4,6 +4,10 @@ import { type CoreUserMessage, generateText } from 'ai';
 import { cookies } from 'next/headers';
 
 import { customModel } from '@/lib/ai';
+import {
+  deleteMessagesByChatIdAfterTimestamp,
+  getMessageById,
+} from '@/lib/db/queries';
 
 export async function saveModelId(model: string) {
   const cookieStore = await cookies();
@@ -27,3 +31,12 @@ export async function generateTitleFromUserMessage({
 
   return title;
 }
+
+export async function deleteTrailingMessages({ id }: { id: string }) {
+  const [message] = await getMessageById({ id });
+
+  await deleteMessagesByChatIdAfterTimestamp({
+    chatId: message.chatId,
+    timestamp: message.createdAt,
+  });
+}
diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
@@ -81,14 +81,21 @@ export async function POST(request: Request) {
     await saveChat({ id, userId: session.user.id, title });
   }
 
+  const userMessageId = generateUUID();
+
   await saveMessages({
     messages: [
-      { ...userMessage, id: generateUUID(), createdAt: new Date(), chatId: id },
+      { ...userMessage, id: userMessageId, createdAt: new Date(), chatId: id },
     ],
   });
 
   const streamingData = new StreamData();
 
+  streamingData.append({
+    type: 'user-message-id',
+    content: userMessageId,
+  });
+
   const result = streamText({
     model: customModel(model.apiIdentifier),
     system: systemPrompt,

diff --git a/components/block-messages.tsx b/components/block-messages.tsx
@@ -3,7 +3,7 @@ import { UIBlock } from './block';
 import { PreviewMessage } from './message';
 import { useScrollToBottom } from './use-scroll-to-bottom';
 import { Vote } from '@/lib/db/schema';
-import { Message } from 'ai';
+import { ChatRequestOptions, Message } from 'ai';
 
 interface BlockMessagesProps {
   chatId: string;
@@ -12,6 +12,12 @@ interface BlockMessagesProps {
   isLoading: boolean;
   votes: Array<Vote> | undefined;
   messages: Array<Message>;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
 }
 
 function PureBlockMessages({
@@ -21,6 +27,8 @@ function PureBlockMessages({
   isLoading,
   votes,
   messages,
+  setMessages,
+  reload,
 }: BlockMessagesProps) {
   const [messagesContainerRef, messagesEndRef] =
     useScrollToBottom<HTMLDivElement>();
@@ -43,6 +51,8 @@ function PureBlockMessages({
               ? votes.find((vote) => vote.messageId === message.id)
               : undefined
           }
+          setMessages={setMessages}
+          reload={reload}
         />
       ))}
 

diff --git a/components/block.tsx b/components/block.tsx
@@ -58,6 +58,7 @@ function PureBlock({
   setBlock,
   messages,
   setMessages,
+  reload,
   votes,
 }: {
   chatId: string;
@@ -82,6 +83,9 @@ function PureBlock({
     },
     chatRequestOptions?: ChatRequestOptions,
   ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
 }) {
   const {
     data: documents,
@@ -286,6 +290,8 @@ function PureBlock({
               setBlock={setBlock}
               votes={votes}
               messages={messages}
+              setMessages={setMessages}
+              reload={reload}
             />
 
             <form className="flex flex-row gap-2 relative items-end w-full px-4 pb-4">

diff --git a/components/chat.tsx b/components/chat.tsx
@@ -36,8 +36,10 @@ export function Chat({
     append,
     isLoading,
     stop,
+    reload,
     data: streamingData,
   } = useChat({
+    id,
     body: { id, modelId: selectedModelId },
     initialMessages,
     onFinish: () => {
@@ -81,6 +83,8 @@ export function Chat({
           isLoading={isLoading}
           votes={votes}
           messages={messages}
+          setMessages={setMessages}
+          reload={reload}
         />
 
         <form className="flex mx-auto px-4 bg-background pb-4 md:pb-6 gap-2 w-full md:max-w-3xl">
@@ -116,6 +120,7 @@ export function Chat({
             setBlock={setBlock}
             messages={messages}
             setMessages={setMessages}
+            reload={reload}
             votes={votes}
           />
         )}

diff --git a/components/message-editor.tsx b/components/message-editor.tsx
@@ -0,0 +1,113 @@
+'use client';
+
+import { ChatRequestOptions, Message } from 'ai';
+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 { useUserMessageId } from '@/hooks/use-user-message-id';
+
+export type MessageEditorProps = {
+  message: Message;
+  setMode: Dispatch<SetStateAction<'view' | 'edit'>>;
+  setMessages: (
+    messages: Message[] | ((messages: Message[]) => Message[]),
+  ) => void;
+  reload: (
+    chatRequestOptions?: ChatRequestOptions,
+  ) => Promise<string | null | undefined>;
+};
+
+export function MessageEditor({
+  message,
+  setMode,
+  setMessages,
+  reload,
+}: MessageEditorProps) {
+  const { userMessageIdFromServer } = useUserMessageId();
+  const [isSubmitting, setIsSubmitting] = useState<boolean>(false);
+
+  const [draftContent, setDraftContent] = useState<string>(message.content);
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
+
+  useEffect(() => {
+    if (textareaRef.current) {
+      adjustHeight();
+    }
+  }, []);
+
+  const adjustHeight = () => {
+    if (textareaRef.current) {
+      textareaRef.current.style.height = 'auto';
+      textareaRef.current.style.height = `${textareaRef.current.scrollHeight + 2}px`;
+    }
+  };
+
+  const handleInput = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
+    setDraftContent(event.target.value);
+    adjustHeight();
+  };
+
+  return (
+    <div className="flex flex-col gap-2 w-full">
+      <Textarea
+        ref={textareaRef}
+        className="bg-transparent outline-none overflow-hidden resize-none !text-base rounded-xl w-full"
+        value={draftContent}
+        onChange={handleInput}
+      />
+
+      <div className="flex flex-row gap-2 justify-end">
+        <Button
+          variant="outline"
+          className="h-fit py-2 px-3"
+          onClick={() => {
+            setMode('view');
+          }}
+        >
+          Cancel
+        </Button>
+        <Button
+          variant="default"
+          className="h-fit py-2 px-3"
+          disabled={isSubmitting}
+          onClick={async () => {
+            setIsSubmitting(true);
+            const messageId = userMessageIdFromServer ?? message.id;
+
+            if (!messageId) {
+              toast.error('Something went wrong, please try again!');
+              setIsSubmitting(false);
+              return;
+            }
+
+            await deleteTrailingMessages({
+              id: messageId,
+            });
+
+            setMessages((messages) => {
+              const index = messages.findIndex((m) => m.id === message.id);
+
+              if (index !== -1) {
+                const updatedMessage = {
+                  ...message,
+                  content: draftContent,
+                };
+
+                return [...messages.slice(0, index), updatedMessage];
+              }
+
+              return messages;
+            });
+
+            setMode('view');
+            reload();
+          }}
+        >
+          {isSubmitting ? 'Sending...' : 'Send'}
+        </Button>
+      </div>
+    </div>
+  );
+}