chore: tweak migration script to handle message order (#969)

Author: jeremyphilemon

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

@@ -3,13 +3,14 @@ import postgres from 'postgres';
 import {
   chat,
   message,
+  type MessageDeprecated,
   messageDeprecated,
   vote,
   voteDeprecated,
 } from '../schema';
 import { drizzle } from 'drizzle-orm/postgres-js';
 import { inArray } from 'drizzle-orm';
-import { appendResponseMessages, UIMessage } from 'ai';
+import { appendResponseMessages, type UIMessage } from 'ai';
 
 config({
   path: '.env.local',
@@ -22,8 +23,8 @@ if (!process.env.POSTGRES_URL) {
 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
+const BATCH_SIZE = 100; // Process 100 chats at a time
+const INSERT_BATCH_SIZE = 1000; // Insert 1000 messages at a time
 
 type NewMessageInsert = {
   id: string;
@@ -40,16 +41,66 @@ type NewVoteInsert = {
   isUpvoted: boolean;
 };
 
-async function createNewTable() {
+interface MessageDeprecatedContentPart {
+  type: string;
+  content: unknown;
+}
+
+function getMessageRank(message: MessageDeprecated): number {
+  if (
+    message.role === 'assistant' &&
+    (message.content as MessageDeprecatedContentPart[]).some(
+      (contentPart) => contentPart.type === 'tool-call',
+    )
+  ) {
+    return 0;
+  }
+
+  if (
+    message.role === 'tool' &&
+    (message.content as MessageDeprecatedContentPart[]).some(
+      (contentPart) => contentPart.type === 'tool-result',
+    )
+  ) {
+    return 1;
+  }
+
+  if (message.role === 'assistant') {
+    return 2;
+  }
+
+  return 3;
+}
+
+function dedupeParts<T extends { type: string; [k: string]: any }>(
+  parts: T[],
+): T[] {
+  const seen = new Set<string>();
+  return parts.filter((p) => {
+    const key = `${p.type}|${JSON.stringify(p.content ?? p)}`;
+    if (seen.has(key)) return false;
+    seen.add(key);
+    return true;
+  });
+}
+
+function sanitizeParts<T extends { type: string; [k: string]: any }>(
+  parts: T[],
+): T[] {
+  return parts.filter(
+    (part) => !(part.type === 'reasoning' && part.reasoning === 'undefined'),
+  );
+}
+
+async function migrateMessages() {
   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)
@@ -60,20 +111,25 @@ async function createNewTable() {
       .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 messages = allMessages
+        .filter((message) => message.chatId === chat.id)
+        .sort((a, b) => {
+          const differenceInTime =
+            new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime();
+          if (differenceInTime !== 0) return differenceInTime;
+
+          return getMessageRank(a) - getMessageRank(b);
+        });
+
       const votes = allVotes.filter((v) => v.chatId === chat.id);
 
-      // Group messages into sections
       const messageSection: Array<UIMessage> = [];
       const messageSections: Array<Array<UIMessage>> = [];
 
@@ -93,7 +149,6 @@ async function createNewTable() {
         messageSections.push([...messageSection]);
       }
 
-      // Process each message section
       for (const section of messageSections) {
         const [userMessage, ...assistantMessages] = section;
 
@@ -121,10 +176,14 @@ async function createNewTable() {
                   attachments: [],
                 } as NewMessageInsert;
               } else if (message.role === 'assistant') {
+                const cleanParts = sanitizeParts(
+                  dedupeParts(message.parts || []),
+                );
+
                 return {
                   id: message.id,
                   chatId: chat.id,
-                  parts: message.parts || [],
+                  parts: cleanParts,
                   role: message.role,
                   createdAt: message.createdAt,
                   attachments: [],
@@ -134,7 +193,6 @@ async function createNewTable() {
             })
             .filter((msg): msg is NewMessageInsert => msg !== null);
 
-          // Add messages to batch
           for (const msg of projectedUISection) {
             newMessagesToInsert.push(msg);
 
@@ -155,11 +213,9 @@ async function createNewTable() {
       }
     }
 
-    // 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,
@@ -173,7 +229,6 @@ async function createNewTable() {
       }
     }
 
-    // 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) {
@@ -185,7 +240,7 @@ async function createNewTable() {
   console.info(`Migration completed: ${processedCount} chats processed`);
 }
 
-createNewTable()
+migrateMessages()
   .then(() => {
     console.info('Script completed successfully');
     process.exit(0);

package.json

@@ -1,6 +1,6 @@
 {
   "name": "ai-chatbot",
-  "version": "3.0.13",
+  "version": "3.0.14",
   "private": true,
   "scripts": {
     "dev": "next dev --turbo",