[Mongo] count messages (aggregate) only when needed V2 (#866)

fixes #863

needed
https://github.com/huggingface/chat-ui/pull/866/commits/9e67e76d75d7257abc5573ae3ec01a8a0040a784

---------

Co-authored-by: Eliott C. <[email protected]>

src/routes/+layout.server.ts

@@ -46,26 +46,6 @@ export const load: LayoutServerLoad = async ({ locals, depends }) => {
 		});
 	}
 
-	// get the number of messages where `from === "assistant"` across all conversations.
-	const totalMessages =
-		(
-			await collections.conversations
-				.aggregate([
-					{ $match: authCondition(locals) },
-					{ $project: { messages: 1 } },
-					{ $unwind: "$messages" },
-					{ $match: { "messages.from": "assistant" } },
-					{ $count: "messages" },
-				])
-				.toArray()
-		)[0]?.messages ?? 0;
-
-	const messagesBeforeLogin = MESSAGES_BEFORE_LOGIN ? parseInt(MESSAGES_BEFORE_LOGIN) : 0;
-
-	const userHasExceededMessages = messagesBeforeLogin > 0 && totalMessages > messagesBeforeLogin;
-
-	const loginRequired = requiresUser && !locals.user && userHasExceededMessages;
-
 	const enableAssistants = ENABLE_ASSISTANTS === "true";
 
 	const assistantActive = !models.map(({ id }) => id).includes(settings?.activeModel ?? "");
@@ -103,6 +83,33 @@ export const load: LayoutServerLoad = async ({ locals, depends }) => {
 
 	const assistants = await collections.assistants.find({ _id: { $in: assistantIds } }).toArray();
 
+	const messagesBeforeLogin = MESSAGES_BEFORE_LOGIN ? parseInt(MESSAGES_BEFORE_LOGIN) : 0;
+
+	let loginRequired = false;
+
+	if (requiresUser && !locals.user && messagesBeforeLogin) {
+		if (conversations.length > messagesBeforeLogin) {
+			loginRequired = true;
+		} else {
+			// get the number of messages where `from === "assistant"` across all conversations.
+			const totalMessages =
+				(
+					await collections.conversations
+						.aggregate([
+							{ $match: { ...authCondition(locals), "messages.from": "assistant" } },
+							{ $project: { messages: 1 } },
+							{ $limit: messagesBeforeLogin + 1 },
+							{ $unwind: "$messages" },
+							{ $match: { "messages.from": "assistant" } },
+							{ $count: "messages" },
+						])
+						.toArray()
+				)[0]?.messages ?? 0;
+
+			loginRequired = totalMessages > messagesBeforeLogin;
+		}
+	}
+
 	return {
 		conversations: conversations.map((conv) => {
 			if (settings?.hideEmojiOnSidebar) {

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

@@ -72,18 +72,17 @@ export async function POST({ request, locals, params, getClientAddress }) {
 		ip: getClientAddress(),
 	});
 
+	const messagesBeforeLogin = MESSAGES_BEFORE_LOGIN ? parseInt(MESSAGES_BEFORE_LOGIN) : 0;
+
 	// guest mode check
-	if (
-		!locals.user?._id &&
-		requiresUser &&
-		(MESSAGES_BEFORE_LOGIN ? parseInt(MESSAGES_BEFORE_LOGIN) : 0) > 0
-	) {
+	if (!locals.user?._id && requiresUser && messagesBeforeLogin) {
 		const totalMessages =
 			(
 				await collections.conversations
 					.aggregate([
-						{ $match: authCondition(locals) },
+						{ $match: { ...authCondition(locals), "messages.from": "assistant" } },
 						{ $project: { messages: 1 } },
+						{ $limit: messagesBeforeLogin + 1 },
 						{ $unwind: "$messages" },
 						{ $match: { "messages.from": "assistant" } },
 						{ $count: "messages" },
@@ -91,7 +90,7 @@ export async function POST({ request, locals, params, getClientAddress }) {
 					.toArray()
 			)[0]?.messages ?? 0;
 
-		if (totalMessages > parseInt(MESSAGES_BEFORE_LOGIN)) {
+		if (totalMessages > messagesBeforeLogin) {
 			throw error(429, "Exceeded number of messages before login");
 		}
 	}