Reduce message update size returned by websearch (#1179)

* Reduce message update size returned by websearch

* lint

* Update src/lib/migrations/routines/06-trim-message-updates.ts

Co-authored-by: Liam Dyer <liamcdyer@gmail.com>

* clean up types

* clean up migrations

---------

Co-authored-by: Liam Dyer <liamcdyer@gmail.com>

src/lib/migrations/routines/04-update-message-updates.ts

@@ -171,7 +171,6 @@ const updateMessageUpdates: Migration = {
 					const webSearchFinishedUpdate: MessageWebSearchFinishedUpdate = {
 						type: MessageUpdateType.WebSearch,
 						subtype: MessageWebSearchUpdateType.Finished,
-						webSearch: message.webSearch,
 					};
 					updates.splice(webSearchSourcesUpdateIndex + 1, 0, webSearchFinishedUpdate);
 				}

src/lib/migrations/routines/06-trim-message-updates.ts

@@ -0,0 +1,68 @@
+import type { Migration } from ".";
+import { collections } from "$lib/server/database";
+import { ObjectId, type WithId } from "mongodb";
+import type { Conversation } from "$lib/types/Conversation";
+import {
+	MessageUpdateType,
+	MessageWebSearchUpdateType,
+	type MessageUpdate,
+} from "$lib/types/MessageUpdate";
+import type { Message } from "$lib/types/Message";
+
+// -----------
+
+/** Converts the old message update to the new schema */
+function convertMessageUpdate(message: Message, update: MessageUpdate): MessageUpdate | null {
+	try {
+		// trim final websearch update, and sources update
+
+		if (update.type === "webSearch") {
+			if (update.subtype === MessageWebSearchUpdateType.Sources) {
+				return {
+					type: MessageUpdateType.WebSearch,
+					subtype: MessageWebSearchUpdateType.Sources,
+					message: update.message,
+					sources: update.sources.map(({ link, title }) => ({ link, title })),
+				};
+			} else if (update.subtype === MessageWebSearchUpdateType.Finished) {
+				return {
+					type: MessageUpdateType.WebSearch,
+					subtype: MessageWebSearchUpdateType.Finished,
+				};
+			}
+		}
+
+		return update;
+	} catch (error) {
+		console.error("Error converting message update during migration. Skipping it... Error:", error);
+		return null;
+	}
+}
+
+const trimMessageUpdates: Migration = {
+	_id: new ObjectId(6),
+	name: "Trim message updates to reduce stored size",
+	up: async () => {
+		const allConversations = collections.conversations.find({});
+
+		let conversation: WithId<Pick<Conversation, "messages">> | null = null;
+		while ((conversation = await allConversations.tryNext())) {
+			const messages = conversation.messages.map((message) => {
+				// Convert all of the existing updates to the new schema
+				const updates = message.updates
+					?.map((update) => convertMessageUpdate(message, update))
+					.filter((update): update is MessageUpdate => Boolean(update));
+
+				return { ...message, updates };
+			});
+
+			// Set the new messages array
+			await collections.conversations.updateOne({ _id: conversation._id }, { $set: { messages } });
+		}
+
+		return true;
+	},
+	runEveryTime: false,
+};
+
+export default trimMessageUpdates;

src/lib/migrations/routines/index.ts

@@ -6,6 +6,7 @@ import type { Database } from "$lib/server/database";
 import addToolsToSettings from "./03-add-tools-in-settings";
 import updateMessageUpdates from "./04-update-message-updates";
 import updateMessageFiles from "./05-update-message-files";
+import trimMessageUpdates from "./06-trim-message-updates";
 
 export interface Migration {
 	_id: ObjectId;
@@ -23,4 +24,5 @@ export const migrations: Migration[] = [
 	addToolsToSettings,
 	updateMessageUpdates,
 	updateMessageFiles,
+	trimMessageUpdates,
 ];

src/lib/server/websearch/runWebSearch.ts

@@ -82,7 +82,7 @@ export async function* runWebSearch(
 			createdAt,
 			updatedAt,
 		};
-		yield makeFinalAnswerUpdate(webSearch);
+		yield makeFinalAnswerUpdate();
 		return webSearch;
 	} catch (searchError) {
 		const message = searchError instanceof Error ? searchError.message : String(searchError);
@@ -97,7 +97,7 @@ export async function* runWebSearch(
 			createdAt,
 			updatedAt,
 		};
-		yield makeFinalAnswerUpdate(webSearch);
+		yield makeFinalAnswerUpdate();
 		return webSearch;
 	}
 }

src/lib/server/websearch/update.ts

@@ -1,4 +1,4 @@
-import type { WebSearch, WebSearchSource } from "$lib/types/WebSearch";
+import type { WebSearchSource } from "$lib/types/WebSearch";
 import {
 	MessageUpdateType,
 	MessageWebSearchUpdateType,
@@ -33,14 +33,13 @@ export function makeSourcesUpdate(sources: WebSearchSource[]): MessageWebSearchS
 		type: MessageUpdateType.WebSearch,
 		subtype: MessageWebSearchUpdateType.Sources,
 		message: "sources",
-		sources,
+		sources: sources.map(({ link, title }) => ({ link, title })),
 	};
 }
 
-export function makeFinalAnswerUpdate(webSearch: WebSearch): MessageWebSearchFinishedUpdate {
+export function makeFinalAnswerUpdate(): MessageWebSearchFinishedUpdate {
 	return {
 		type: MessageUpdateType.WebSearch,
 		subtype: MessageWebSearchUpdateType.Finished,
-		webSearch,
 	};
 }

src/lib/types/MessageUpdate.ts

@@ -1,4 +1,4 @@
-import type { WebSearch, WebSearchSource } from "$lib/types/WebSearch";
+import type { WebSearchSource } from "$lib/types/WebSearch";
 import type { ToolCall, ToolResult } from "$lib/types/Tool";
 
 export type MessageUpdate =
@@ -39,7 +39,7 @@ export enum MessageWebSearchUpdateType {
 	Sources = "sources",
 	Finished = "finished",
 }
-interface BaseMessageWebSearchUpdate<TSubType extends MessageWebSearchUpdateType> {
+export interface BaseMessageWebSearchUpdate<TSubType extends MessageWebSearchUpdateType> {
 	type: MessageUpdateType.WebSearch;
 	subtype: TSubType;
 }
@@ -58,10 +58,8 @@ export interface MessageWebSearchSourcesUpdate
 	message: string;
 	sources: WebSearchSource[];
 }
-export interface MessageWebSearchFinishedUpdate
-	extends BaseMessageWebSearchUpdate<MessageWebSearchUpdateType.Finished> {
-	webSearch: WebSearch;
-}
+export type MessageWebSearchFinishedUpdate =
+	BaseMessageWebSearchUpdate<MessageWebSearchUpdateType.Finished>;
 export type MessageWebSearchUpdate =
 	| MessageWebSearchErrorUpdate
 	| MessageWebSearchGeneralUpdate

src/routes/conversation/[id]/+server.ts

@@ -11,7 +11,6 @@ import { z } from "zod";
 import {
 	MessageUpdateStatus,
 	MessageUpdateType,
-	MessageWebSearchUpdateType,
 	type MessageUpdate,
 } from "$lib/types/MessageUpdate";
 import { uploadFile } from "$lib/server/files/uploadFile";
@@ -359,14 +358,6 @@ export async function POST({ request, locals, params, getClientAddress }) {
 					];
 				}
 
-				// Set web search
-				else if (
-					event.type === MessageUpdateType.WebSearch &&
-					event.subtype === MessageWebSearchUpdateType.Finished
-				) {
-					messageToWriteTo.webSearch = event.webSearch;
-				}
-
 				// Append to the persistent message updates if it's not a stream update
 				if (event.type !== "stream") {
 					messageToWriteTo?.updates?.push(event);