Merge branch 'main' of github.com:athrael-soju/ai-chatbot

athrael-soju · Jan 28, 2025 · 6e3dbb0 · 6e3dbb0
2 parents 7b59b5e + 5bb62f1
commit 6e3dbb0
Show file tree

Hide file tree

Showing 9 changed files with 41 additions and 79 deletions.
diff --git a/app/(chat)/actions.ts b/app/(chat)/actions.ts
@@ -1,6 +1,6 @@
 'use server';
 
-import { type CoreUserMessage, generateText } from 'ai';
+import { type CoreUserMessage, generateText, Message } from 'ai';
 import { cookies } from 'next/headers';
 
 import { customModel } from '@/lib/ai';
@@ -19,7 +19,7 @@ export async function saveModelId(model: string) {
 export async function generateTitleFromUserMessage({
   message,
 }: {
-  message: CoreUserMessage;
+  message: Message;
 }) {
   const { text: title } = await generateText({
     model: customModel('gpt-4o-mini'),

diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts
@@ -1,6 +1,5 @@
 import {
   type Message,
-  convertToCoreMessages,
   createDataStreamResponse,
   smoothStream,
   streamText,
@@ -64,9 +63,8 @@ export async function POST(request: Request) {
   if (!model) {
     return new Response('Model not found', { status: 404 });
   }
-  // TODO: Access knowledgebase here and enhance the user message.
-  const coreMessages = convertToCoreMessages(messages);
-  const userMessage = getMostRecentUserMessage(coreMessages);
+
+  const userMessage = getMostRecentUserMessage(messages);
 
   if (!userMessage) {
     return new Response('No user message found', { status: 400 });
@@ -79,28 +77,20 @@ export async function POST(request: Request) {
     await saveChat({ id, userId: session.user.id, title });
   }
 
-  const userMessageId = generateUUID();
-  // TODO: Save original user message here.
   await saveMessages({
-    messages: [
-      { ...userMessage, id: userMessageId, createdAt: new Date(), chatId: id },
-    ],
+    messages: [{ ...userMessage, createdAt: new Date(), chatId: id }],
   });
 
   return createDataStreamResponse({
     execute: (dataStream) => {
-      dataStream.writeData({
-        type: 'user-message-id',
-        content: userMessageId,
-      });
-
       const result = streamText({
         model: customModel(model.apiIdentifier),
         system: systemPrompt,
-        messages: coreMessages,
+        messages,
         maxSteps: 5,
         experimental_activeTools: allTools,
         experimental_transform: smoothStream({ chunking: 'word' }),
+        experimental_generateMessageId: generateUUID,
         tools: {
           getWeather,
           createDocument: createDocument({ session, dataStream, model }),
@@ -120,16 +110,8 @@ export async function POST(request: Request) {
               await saveMessages({
                 messages: responseMessagesWithoutIncompleteToolCalls.map(
                   (message) => {
-                    const messageId = generateUUID();
-
-                    if (message.role === 'assistant') {
-                      dataStream.writeMessageAnnotation({
-                        messageIdFromServer: messageId,
-                      });
-                    }
-
                     return {
-                      id: messageId,
+                      id: message.id,
                       chatId: id,
                       role: message.role,
                       content: message.content,

diff --git a/components/chat.tsx b/components/chat.tsx
@@ -7,7 +7,7 @@ import useSWR, { useSWRConfig } from 'swr';
 
 import { ChatHeader } from '@/components/chat-header';
 import type { Vote } from '@/lib/db/schema';
-import { fetcher } from '@/lib/utils';
+import { fetcher, generateUUID } from '@/lib/utils';
 
 import { Block } from './block';
 import { MultimodalInput } from './multimodal-input';
@@ -45,6 +45,8 @@ export function Chat({
     body: { id, modelId: selectedModelId },
     initialMessages,
     experimental_throttle: 100,
+    sendExtraMessageFields: true,
+    generateId: generateUUID,
     onFinish: () => {
       mutate('/api/history');
     },

diff --git a/components/data-stream-handler.tsx b/components/data-stream-handler.tsx
@@ -5,7 +5,6 @@ import { useEffect, useRef } from 'react';
 import { blockDefinitions, BlockKind } from './block';
 import { Suggestion } from '@/lib/db/schema';
 import { initialBlockData, useBlock } from '@/hooks/use-block';
-import { useUserMessageId } from '@/hooks/use-user-message-id';
 
 export type DataStreamDelta = {
   type:
@@ -17,14 +16,12 @@ export type DataStreamDelta = {
     | 'suggestion'
     | 'clear'
     | 'finish'
-    | 'user-message-id'
     | 'kind';
   content: string | Suggestion;
 };
 
 export function DataStreamHandler({ id }: { id: string }) {
   const { data: dataStream } = useChat({ id });
-  const { setUserMessageIdFromServer } = useUserMessageId();
   const { block, setBlock, setMetadata } = useBlock();
   const lastProcessedIndex = useRef(-1);
 
@@ -35,11 +32,6 @@ export function DataStreamHandler({ id }: { id: string }) {
     lastProcessedIndex.current = dataStream.length - 1;
 
     (newDeltas as DataStreamDelta[]).forEach((delta: DataStreamDelta) => {
-      if (delta.type === 'user-message-id') {
-        setUserMessageIdFromServer(delta.content as string);
-        return;
-      }
-
       const blockDefinition = blockDefinitions.find(
         (blockDefinition) => blockDefinition.kind === block.kind,
       );
@@ -97,7 +89,7 @@ export function DataStreamHandler({ id }: { id: string }) {
         }
       });
     });
-  }, [dataStream, setBlock, setUserMessageIdFromServer, setMetadata, block]);
+  }, [dataStream, setBlock, setMetadata, block]);
 
   return null;
 }
diff --git a/components/message-actions.tsx b/components/message-actions.tsx
@@ -4,7 +4,6 @@ import { useSWRConfig } from 'swr';
 import { useCopyToClipboard } from 'usehooks-ts';
 
 import type { Vote } from '@/lib/db/schema';
-import { getMessageIdFromAnnotations } from '@/lib/utils';
 
 import { CopyIcon, ThumbDownIcon, ThumbUpIcon } from './icons';
 import { Button } from './ui/button';
@@ -62,13 +61,11 @@ export function PureMessageActions({
               disabled={vote?.isUpvoted}
               variant="outline"
               onClick={async () => {
-                const messageId = getMessageIdFromAnnotations(message);
-
                 const upvote = fetch('/api/vote', {
                   method: 'PATCH',
                   body: JSON.stringify({
                     chatId,
-                    messageId,
+                    messageId: message.id,
                     type: 'up',
                   }),
                 });
@@ -116,13 +113,11 @@ export function PureMessageActions({
               variant="outline"
               disabled={vote && !vote.isUpvoted}
               onClick={async () => {
-                const messageId = getMessageIdFromAnnotations(message);
-
                 const downvote = fetch('/api/vote', {
                   method: 'PATCH',
                   body: JSON.stringify({
                     chatId,
-                    messageId,
+                    messageId: message.id,
                     type: 'down',
                   }),
                 });

diff --git a/components/message-editor.tsx b/components/message-editor.tsx
@@ -6,7 +6,6 @@ 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;
@@ -25,7 +24,6 @@ export function MessageEditor({
   setMessages,
   reload,
 }: MessageEditorProps) {
-  const { userMessageIdFromServer } = useUserMessageId();
   const [isSubmitting, setIsSubmitting] = useState<boolean>(false);
 
   const [draftContent, setDraftContent] = useState<string>(message.content);
@@ -74,16 +72,9 @@ export function MessageEditor({
           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,
+              id: message.id,
             });
 
             setMessages((messages) => {

diff --git a/hooks/use-user-message-id.ts b/hooks/use-user-message-id.ts
diff --git a/lib/db/queries.ts b/lib/db/queries.ts
@@ -1,7 +1,7 @@
 import 'server-only';
 
 import { genSaltSync, hashSync } from 'bcrypt-ts';
-import { and, asc, desc, eq, gt, gte } from 'drizzle-orm';
+import { and, asc, desc, eq, gt, gte, inArray } from 'drizzle-orm';
 import { drizzle } from 'drizzle-orm/postgres-js';
 import postgres from 'postgres';
 
@@ -301,11 +301,28 @@ export async function deleteMessagesByChatIdAfterTimestamp({
   timestamp: Date;
 }) {
   try {
-    return await db
-      .delete(message)
+    const messagesToDelete = await db
+      .select({ id: message.id })
+      .from(message)
       .where(
         and(eq(message.chatId, chatId), gte(message.createdAt, timestamp)),
       );
+
+    const messageIds = messagesToDelete.map((message) => message.id);
+
+    if (messageIds.length > 0) {
+      await db
+        .delete(vote)
+        .where(
+          and(eq(vote.chatId, chatId), inArray(vote.messageId, messageIds)),
+        );
+
+      return await db
+        .delete(message)
+        .where(
+          and(eq(message.chatId, chatId), inArray(message.id, messageIds)),
+        );
+    }
   } catch (error) {
     console.error(
       'Failed to delete messages by id after timestamp from database',

diff --git a/lib/utils.ts b/lib/utils.ts
@@ -128,9 +128,12 @@ export function convertToUIMessages(
   }, []);
 }
 
+type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage;
+type ResponseMessage = ResponseMessageWithoutId & { id: string };
+
 export function sanitizeResponseMessages(
-  messages: Array<CoreToolMessage | CoreAssistantMessage>,
-): Array<CoreToolMessage | CoreAssistantMessage> {
+  messages: Array<ResponseMessage>,
+): Array<ResponseMessage> {
   const toolResultIds: Array<string> = [];
 
   for (const message of messages) {
@@ -200,7 +203,7 @@ export function sanitizeUIMessages(messages: Array<Message>): Array<Message> {
   );
 }
 
-export function getMostRecentUserMessage(messages: Array<CoreMessage>) {
+export function getMostRecentUserMessage(messages: Array<Message>) {
   const userMessages = messages.filter((message) => message.role === 'user');
   return userMessages.at(-1);
 }
@@ -214,13 +217,3 @@ export function getDocumentTimestampByIndex(
 
   return documents[index].createdAt;
 }
-
-export function getMessageIdFromAnnotations(message: Message) {
-  if (!message.annotations) return message.id;
-
-  const [annotation] = message.annotations;
-  if (!annotation) return message.id;
-
-  // @ts-expect-error messageIdFromServer is not defined in MessageAnnotation
-  return annotation.messageIdFromServer;
-}