2 files changed, 73 insertions, 2 deletions

diff --git a/apps/web/app/api/chat/route.ts b/apps/web/app/api/chat/route.ts
index 004bfd3b..a14c96df 100644
--- a/apps/web/app/api/chat/route.ts
+++ b/apps/web/app/api/chat/route.ts
@@ -7,6 +7,10 @@ import {
 } from "@repo/shared-types";
 import { ensureAuth } from "../ensureAuth";
 import { z } from "zod";
+import { db } from "@/server/db";
+import { chatHistory as chatHistoryDb, chatThreads } from "@/server/db/schema";
+import { and, eq, gt, sql } from "drizzle-orm";
+import { join } from "path";
 
 export const runtime = "edge";
 
@@ -21,12 +25,56 @@ export async function POST(req: NextRequest) {
 		return new Response("Missing BACKEND_SECURITY_KEY", { status: 500 });
 	}
 
+	const ip = req.headers.get("cf-connecting-ip");
+
+	if (ip) {
+		if (process.env.RATELIMITER) {
+			const { success } = await process.env.RATELIMITER.limit({
+				key: `chat-${ip}`,
+			});
+
+			if (!success) {
+				console.error("rate limit exceeded");
+				return new Response("Rate limit exceeded", { status: 429 });
+			}
+		} else {
+			console.info("RATELIMITER not found in env");
+		}
+	} else {
+		console.info("cf-connecting-ip not found in headers");
+	}
+
+	const lastHour = new Date(new Date().getTime() - 3600000);
+
+	// Only allow 5 requests per hour for each user, something lke this but this one is bad because chathistory.userid doesnt exist, we have to do a join and get it from the threads table
+	const result = await db
+		.select({
+			count: sql<number>`count(*)`.mapWith(Number),
+		})
+		.from(chatHistoryDb)
+		.innerJoin(chatThreads, eq(chatHistoryDb.threadId, chatThreads.id))
+		.where(
+			and(
+				eq(chatThreads.userId, session.user.id),
+				gt(chatHistoryDb.createdAt, lastHour),
+			),
+		)
+		.execute();
+
+	if (result[0]?.count && result[0]?.count >= 5) {
+		// return new Response(`Too many requests ${result[0]?.count}`, { status: 429 });
+		console.log(result[0]?.count);
+	} else {
+		console.log("count", result);
+	}
+
 	const url = new URL(req.url);
 
 	const query = url.searchParams.get("q");
 	const spaces = url.searchParams.get("spaces");
 
 	const sourcesOnly = url.searchParams.get("sourcesOnly") ?? "false";
+	const proMode = url.searchParams.get("proMode") === "true";
 
 	const jsonRequest = (await req.json()) as {
 		chatHistory: ChatHistory[];
@@ -55,7 +103,7 @@ export async function POST(req: NextRequest) {
 	const modelCompatible = await convertChatHistoryList(validated.data);
 
 	const resp = await fetch(
-		`${process.env.BACKEND_BASE_URL}/api/chat?query=${query}&user=${session.user.id}&sourcesOnly=${sourcesOnly}&spaces=${spaces}`,
+		`${process.env.BACKEND_BASE_URL}/api/chat?query=${query}&user=${session.user.id}&sourcesOnly=${sourcesOnly}&spaces=${spaces}&proMode=${proMode}`,
 		{
 			headers: {
 				Authorization: `Bearer ${process.env.BACKEND_SECURITY_KEY}`,
diff --git a/apps/web/app/api/store/route.ts b/apps/web/app/api/store/route.ts
index f9ab7c01..992c2a0e 100644
--- a/apps/web/app/api/store/route.ts
+++ b/apps/web/app/api/store/route.ts
@@ -4,7 +4,7 @@ import { ensureAuth } from "../ensureAuth";
 import { z } from "zod";
 import { db } from "@/server/db";
 import { contentToSpace, space, storedContent } from "@/server/db/schema";
-import { and, eq, inArray } from "drizzle-orm";
+import { and, eq, gt, inArray, sql } from "drizzle-orm";
 import { LIMITS } from "@/lib/constants";
 import { limit } from "@/app/actions/doers";
 
@@ -22,6 +22,29 @@ const createMemoryFromAPI = async (input: {
 		};
 	}
 
+	// Get number of items saved in the last 2 hours
+	const last2Hours = new Date(Date.now() - 2 * 60 * 60 * 1000);
+
+	const numberOfItemsSavedInLast2Hours = await db
+		.select({
+			count: sql<number>`count(*)`.mapWith(Number),
+		})
+		.from(storedContent)
+		.where(
+			and(
+				gt(storedContent.savedAt, last2Hours),
+				eq(storedContent.userId, input.userId),
+			),
+		);
+
+	if (numberOfItemsSavedInLast2Hours[0]!.count >= 20) {
+		return {
+			success: false,
+			data: 0,
+			error: `You have exceeded the limit`,
+		};
+	}
+
 	const vectorSaveResponse = await fetch(
 		`${process.env.BACKEND_BASE_URL}/api/add`,
 		{