adding rate limit to APIs

package.json

@@ -21,6 +21,7 @@
     "@types/node": "20.2.5",
     "@types/react": "18.2.8",
     "@types/react-dom": "18.2.4",
+    "@upstash/ratelimit": "^0.4.3",
     "@upstash/redis": "^1.21.0",
     "ai": "^2.1.3",
     "autoprefixer": "10.4.14",

src/app/api/chatgpt/route.ts

@@ -11,6 +11,7 @@ import { PromptTemplate } from "langchain/prompts";
 import { NextResponse } from "next/server";
 import { currentUser } from "@clerk/nextjs";
 import MemoryManager from "@/app/utils/memory";
+import { rateLimit } from "@/app/utils/rateLimit";
 
 dotenv.config({ path: `.env.local` });
 
@@ -20,6 +21,21 @@ export async function POST(req: Request) {
   let clerkUserName;
   const { prompt, isText, userId, userName } = await req.json();
 
+  const identifier = req.url + "-" + (userId || "anonymous");
+  const { success } = await rateLimit(identifier);
+  if (!success) {
+    console.log("INFO: rate limit exceeded");
+    return new NextResponse(
+      JSON.stringify({ Message: "Hi, the companions can't talk this fast." }),
+      {
+        status: 429,
+        headers: {
+          "Content-Type": "application/json",
+        },
+      }
+    );
+  }
+
   // XXX Companion name passed here. Can use as a key to get backstory, chat history etc.
   const name = req.headers.get("name");
   const companion_file_name = name + ".txt";
@@ -35,6 +51,7 @@ export async function POST(req: Request) {
   }
 
   if (!clerkUserId || !!!(await clerk.users.getUser(clerkUserId))) {
+    console.log("user not authorized");
     return new NextResponse(
       JSON.stringify({ Message: "User not authorized" }),
       {

src/app/api/text/route.ts

@@ -3,11 +3,11 @@ import twilio from "twilio";
 import clerk from "@clerk/clerk-sdk-node";
 import dotenv from "dotenv";
 import ConfigManager from "@/app/utils/config";
+import { rateLimit } from "@/app/utils/rateLimit";
 
 dotenv.config({ path: `.env.local` });
 const twilioAuthToken = process.env.TWILIO_AUTH_TOKEN;
 const accountSid = process.env.TWILIO_ACCOUNT_SID;
-const publishableKey = process.env.NEXT_PUBLIC_CLERK_PUBLISHABLE_KEY;
 
 export async function POST(request: Request) {
   let queryMap: any = {};
@@ -21,6 +21,21 @@ export async function POST(request: Request) {
   const phoneNumber = queryMap["From"];
   const companionPhoneNumber = queryMap["To"];
 
+  const identifier = request.url + "-" + (phoneNumber || "anonymous");
+  const { success } = await rateLimit(identifier);
+  if (!success) {
+    console.log("INFO: rate limit exceeded");
+    return new NextResponse(
+      JSON.stringify({ Message: "Hi, the companions can't talk this fast." }),
+      {
+        status: 429,
+        headers: {
+          "Content-Type": "application/json",
+        },
+      }
+    );
+  }
+
   // check if the user has verified phone #
   const users = await clerk.users.getUserList({ phoneNumber });
 
@@ -37,7 +52,6 @@ export async function POST(request: Request) {
   }
 
   const configManager = ConfigManager.getInstance();
-  console.log(phoneNumber);
   const companionConfig = configManager.getConfig(
     "phone",
     companionPhoneNumber

src/app/api/vicuna13b/route.ts

@@ -9,6 +9,7 @@ import clerk from "@clerk/clerk-sdk-node";
 import MemoryManager from "@/app/utils/memory";
 import { currentUser } from "@clerk/nextjs";
 import { NextResponse } from "next/server";
+import { rateLimit } from "@/app/utils/rateLimit";
 
 dotenv.config({ path: `.env.local` });
 
@@ -17,6 +18,22 @@ export async function POST(request: Request) {
   let clerkUserId;
   let user;
   let clerkUserName;
+
+  const identifier = request.url + "-" + (userId || "anonymous");
+  const { success } = await rateLimit(identifier);
+  if (!success) {
+    console.log("INFO: rate limit exceeded");
+    return new NextResponse(
+      JSON.stringify({ Message: "Hi, the companions can't talk this fast." }),
+      {
+        status: 429,
+        headers: {
+          "Content-Type": "application/json",
+        },
+      }
+    );
+  }
+
   // XXX Companion name passed here. Can use as a key to get backstory, chat history etc.
   const name = request.headers.get("name");
   const companion_file_name = name + ".txt";

src/app/utils/rateLimit.ts

@@ -0,0 +1,13 @@
+import { Ratelimit } from "@upstash/ratelimit";
+import { Redis } from "@upstash/redis";
+
+export async function rateLimit(identifier: string) {
+  // Rate limit through Upstash
+  const ratelimit = new Ratelimit({
+    redis: Redis.fromEnv(),
+    limiter: Ratelimit.slidingWindow(10, "10 s"),
+    analytics: true,
+    prefix: "@upstash/ratelimit",
+  });
+  return await ratelimit.limit(identifier);
+}