From 1fd23029140f3378387934942eb5a8a1e24d2845 Mon Sep 17 00:00:00 2001 From: Jeremy Date: Wed, 30 Apr 2025 16:04:09 -0700 Subject: [PATCH] chore: tweak migration script to handle message order (#969) --- lib/db/helpers/01-core-to-parts.ts | 91 ++++++++++++++++++++++++------ package.json | 2 +- 2 files changed, 74 insertions(+), 19 deletions(-) diff --git a/lib/db/helpers/01-core-to-parts.ts b/lib/db/helpers/01-core-to-parts.ts index 2a6bfbd..81c2693 100644 --- a/lib/db/helpers/01-core-to-parts.ts +++ b/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( + parts: T[], +): T[] { + const seen = new Set(); + 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( + 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 = []; const messageSections: Array> = []; @@ -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); diff --git a/package.json b/package.json index d0a9fd2..f8a2752 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "ai-chatbot", - "version": "3.0.13", + "version": "3.0.14", "private": true, "scripts": { "dev": "next dev --turbo",