Add rate limiting to websearch and title summary (#433)

Author: nsarrazin

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

@@ -1,12 +1,14 @@
+import { RATE_LIMIT } from "$env/static/private";
 import { buildPrompt } from "$lib/buildPrompt";
 import { authCondition } from "$lib/server/auth";
 import { collections } from "$lib/server/database";
 import { generateFromDefaultEndpoint } from "$lib/server/generateFromDefaultEndpoint";
 import { defaultModel } from "$lib/server/models";
+import { ERROR_MESSAGES } from "$lib/stores/errors.js";
 import { error } from "@sveltejs/kit";
 import { ObjectId } from "mongodb";
 
-export async function POST({ params, locals }) {
+export async function POST({ params, locals, getClientAddress }) {
 	const convId = new ObjectId(params.id);
 
 	const conversation = await collections.conversations.findOne({
@@ -18,6 +20,23 @@ export async function POST({ params, locals }) {
 		throw error(404, "Conversation not found");
 	}
 
+	const userId = locals.user?._id ?? locals.sessionId;
+
+	await collections.messageEvents.insertOne({
+		userId: userId,
+		createdAt: new Date(),
+		ip: getClientAddress(),
+	});
+
+	const nEvents = Math.max(
+		await collections.messageEvents.countDocuments({ userId }),
+		await collections.messageEvents.countDocuments({ ip: getClientAddress() })
+	);
+
+	if (RATE_LIMIT != "" && nEvents > parseInt(RATE_LIMIT)) {
+		throw error(429, ERROR_MESSAGES.rateLimited);
+	}
+
 	const firstMessage = conversation.messages.find((m) => m.from === "user");
 
 	const userPrompt =

src/routes/conversation/[id]/web-search/+server.ts

@@ -10,6 +10,8 @@ import type { WebSearch } from "$lib/types/WebSearch";
 import { generateQuery } from "$lib/server/websearch/generateQuery";
 import { parseWeb } from "$lib/server/websearch/parseWeb";
 import { summarizeWeb } from "$lib/server/websearch/summarizeWeb";
+import { RATE_LIMIT } from "$env/static/private";
+import { ERROR_MESSAGES } from "$lib/stores/errors.js";
 
 interface GenericObject {
 	[key: string]: GenericObject | unknown;
@@ -22,7 +24,7 @@ function removeLinks(obj: GenericObject) {
 	}
 	return obj;
 }
-export async function GET({ params, locals, url }) {
+export async function GET({ params, locals, url, getClientAddress }) {
 	const model = defaultModel;
 	const convId = new ObjectId(params.id);
 	const searchId = new ObjectId();
@@ -36,6 +38,23 @@ export async function GET({ params, locals, url }) {
 		throw error(404, "Conversation not found");
 	}
 
+	const userId = locals.user?._id ?? locals.sessionId;
+
+	await collections.messageEvents.insertOne({
+		userId: userId,
+		createdAt: new Date(),
+		ip: getClientAddress(),
+	});
+
+	const nEvents = Math.max(
+		await collections.messageEvents.countDocuments({ userId }),
+		await collections.messageEvents.countDocuments({ ip: getClientAddress() })
+	);
+
+	if (RATE_LIMIT != "" && nEvents > parseInt(RATE_LIMIT)) {
+		throw error(429, ERROR_MESSAGES.rateLimited);
+	}
+
 	const prompt = z.string().trim().min(1).parse(url.searchParams.get("prompt"));
 
 	const messages = (() => {