add rate-limiting

This commit is contained in:
ramnique 2025-02-04 16:35:12 +05:30
parent 024f6c75cc
commit 200e8d2e38
9 changed files with 188 additions and 2 deletions

View file

@ -3,7 +3,7 @@
import { redirect } from "next/navigation";
import { SimulationData, EmbeddingDoc, GetInformationToolResult, DataSource, PlaygroundChat, AgenticAPIChatRequest, AgenticAPIChatResponse, convertFromAgenticAPIChatMessages, WebpageCrawlResponse, Workflow, WorkflowAgent, CopilotAPIRequest, CopilotAPIResponse, CopilotMessage, CopilotWorkflow, convertToCopilotWorkflow, convertToCopilotApiMessage, convertToCopilotMessage, CopilotAssistantMessage, CopilotChatContext, convertToCopilotApiChatContext, Scenario, ClientToolCallRequestBody, ClientToolCallJwt, ClientToolCallRequest, WithStringId, Project, WorkflowTool, WorkflowPrompt, ApiKey } from "./lib/types";
import { ObjectId, WithId } from "mongodb";
import { generateObject, generateText, tool, embed } from "ai";
import { generateObject, generateText, embed } from "ai";
import { dataSourcesCollection, embeddingsCollection, projectsCollection, webpagesCollection, agentWorkflowsCollection, scenariosCollection, projectMembersCollection, apiKeysCollection } from "@/app/lib/mongodb";
import { z } from 'zod';
import { openai } from "@ai-sdk/openai";
@ -12,12 +12,13 @@ import { embeddingModel } from "./lib/embedding";
import { apiV1 } from "rowboat-shared";
import { zodToJsonSchema } from 'zod-to-json-schema';
import crypto from 'crypto';
import { SignJWT } from "jose";
import { Claims, getSession } from "@auth0/nextjs-auth0";
import { revalidatePath } from "next/cache";
import { callClientToolWebhook, getAgenticApiResponse } from "./lib/utils";
import { templates } from "./lib/project_templates";
import { assert, error } from "node:console";
import { check_query_limit } from "./lib/rate_limiting";
import { QueryLimitError } from "./lib/client_utils";
const crawler = new FirecrawlApp({ apiKey: process.env.FIRECRAWL_API_KEY || '' });
@ -319,6 +320,19 @@ export async function scrapeWebpage(url: string): Promise<z.infer<typeof Webpage
export async function createProject(formData: FormData) {
const user = await authCheck();
// ensure that projects created by this user is less than
// configured limit
const projectsLimit = Number(process.env.MAX_PROJECTS_PER_USER) || 0;
if (projectsLimit > 0) {
const count = await projectsCollection.countDocuments({
createdByUserId: user.sub,
});
if (count >= projectsLimit) {
throw new Error('You have reached your project limit. Please upgrade your plan.');
}
}
const name = formData.get('name') as string;
const templateKey = formData.get('template') as string;
const projectId = crypto.randomUUID();
@ -492,6 +506,9 @@ export async function getAssistantResponse(
rawResponse: unknown,
}> {
await projectAuthCheck(projectId);
if (!await check_query_limit(projectId)) {
throw new QueryLimitError();
}
const response = await getAgenticApiResponse(request);
return {
@ -513,6 +530,9 @@ export async function getCopilotResponse(
rawResponse: unknown,
}> {
await projectAuthCheck(projectId);
if (!await check_query_limit(projectId)) {
throw new QueryLimitError();
}
// prepare request
const request: z.infer<typeof CopilotAPIRequest> = {
@ -643,6 +663,9 @@ export async function getCopilotResponse(
export async function suggestToolResponse(toolId: string, projectId: string, messages: z.infer<typeof apiV1.ChatMessage>[]): Promise<string> {
await projectAuthCheck(projectId);
if (!await check_query_limit(projectId)) {
throw new QueryLimitError();
}
const prompt = `
# Your Specific Task:
@ -891,6 +914,10 @@ export async function simulateUserResponse(
simulationData: z.infer<typeof SimulationData>
): Promise<string> {
await projectAuthCheck(projectId);
if (!await check_query_limit(projectId)) {
throw new QueryLimitError();
}
const articlePrompt = `
# Your Specific Task:

View file

@ -5,6 +5,7 @@ import { ObjectId } from "mongodb";
import { authCheck } from "@/app/api/v1/utils";
import { convertFromApiToAgenticApiMessages, convertFromAgenticApiToApiMessages, AgenticAPIChatRequest, ApiRequest, ApiResponse, convertWorkflowToAgenticAPI } from "@/app/lib/types";
import { getAgenticApiResponse } from "@/app/lib/utils";
import { check_query_limit } from "@/app/lib/rate_limiting";
// get next turn / agent response
export async function POST(
@ -13,6 +14,11 @@ export async function POST(
): Promise<Response> {
const { projectId } = await params;
// check query limit
if (!await check_query_limit(projectId)) {
return Response.json({ error: "Query limit exceeded" }, { status: 429 });
}
return await authCheck(projectId, req, async () => {
// parse and validate the request body
let body;

View file

@ -6,6 +6,7 @@ import { ObjectId, WithId } from "mongodb";
import { authCheck } from "../../../utils";
import { AgenticAPIChatRequest, convertFromAgenticAPIChatMessages, convertToAgenticAPIChatMessages, convertWorkflowToAgenticAPI } from "@/app/lib/types";
import { callClientToolWebhook, getAgenticApiResponse } from "@/app/lib/utils";
import { check_query_limit } from "@/app/lib/rate_limiting";
const chatsCollection = db.collection<z.infer<typeof apiV1.Chat>>("chats");
const chatMessagesCollection = db.collection<z.infer<typeof apiV1.ChatMessage>>("chatMessages");
@ -18,6 +19,11 @@ export async function POST(
return await authCheck(req, async (session) => {
const { chatId } = await params;
// check query limit
if (!await check_query_limit(session.projectId)) {
return Response.json({ error: "Query limit exceeded" }, { status: 429 });
}
// parse and validate the request body
let body;
try {

View file

@ -0,0 +1,6 @@
export class QueryLimitError extends Error {
constructor(message: string = 'Query limit exceeded') {
super(message);
this.name = 'QueryLimitError';
}
}

View file

@ -0,0 +1,21 @@
import { redisClient } from "./redis";
const MAX_QUERIES_PER_MINUTE = Number(process.env.MAX_QUERIES_PER_MINUTE) || 0;
export async function check_query_limit(projectId: string): Promise<boolean> {
// if the limit is 0, we don't want to check the limit
if (MAX_QUERIES_PER_MINUTE === 0) {
return true;
}
const minutes_since_epoch = Math.floor(Date.now() / 1000 / 60); // 60 second window
const key = `rate_limit:${projectId}:${minutes_since_epoch}`;
// increment the counter and return the count
const count = await redisClient.incr(key);
if (count === 1) {
await redisClient.expire(key, 70); // Set TTL to clean up automatically
}
return count <= MAX_QUERIES_PER_MINUTE;
}

View file

@ -0,0 +1,7 @@
import { createClient } from 'redis';
export const redisClient = createClient({
url: process.env.REDIS_URL,
});
redisClient.connect();

View file

@ -34,6 +34,7 @@
"react-dom": "^18.3.1",
"react-markdown": "^9.0.1",
"react-resizable-panels": "^2.1.7",
"redis": "^4.7.0",
"remark-gfm": "^4.0.0",
"rowboat-shared": "github:rowboatlabs/shared",
"sharp": "^0.33.4",
@ -50,6 +51,7 @@
"@types/node": "^20",
"@types/react": "^18",
"@types/react-dom": "^18",
"@types/redis": "^4.0.11",
"eslint": "^8",
"eslint-config-next": "14.2.5",
"postcss": "^8",
@ -6440,6 +6442,64 @@
"react": "^16.8.0 || ^17.0.0-rc.1 || ^18.0.0"
}
},
"node_modules/@redis/bloom": {
"version": "1.2.0",
"resolved": "https://registry.npmjs.org/@redis/bloom/-/bloom-1.2.0.tgz",
"integrity": "sha512-HG2DFjYKbpNmVXsa0keLHp/3leGJz1mjh09f2RLGGLQZzSHpkmZWuwJbAvo3QcRY8p80m5+ZdXZdYOSBLlp7Cg==",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/client": {
"version": "1.6.0",
"resolved": "https://registry.npmjs.org/@redis/client/-/client-1.6.0.tgz",
"integrity": "sha512-aR0uffYI700OEEH4gYnitAnv3vzVGXCFvYfdpu/CJKvk4pHfLPEy/JSZyrpQ+15WhXe1yJRXLtfQ84s4mEXnPg==",
"dependencies": {
"cluster-key-slot": "1.1.2",
"generic-pool": "3.9.0",
"yallist": "4.0.0"
},
"engines": {
"node": ">=14"
}
},
"node_modules/@redis/client/node_modules/yallist": {
"version": "4.0.0",
"resolved": "https://registry.npmjs.org/yallist/-/yallist-4.0.0.tgz",
"integrity": "sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A=="
},
"node_modules/@redis/graph": {
"version": "1.1.1",
"resolved": "https://registry.npmjs.org/@redis/graph/-/graph-1.1.1.tgz",
"integrity": "sha512-FEMTcTHZozZciLRl6GiiIB4zGm5z5F3F6a6FZCyrfxdKOhFlGkiAqlexWMBzCi4DcRoyiOsuLfW+cjlGWyExOw==",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/json": {
"version": "1.0.7",
"resolved": "https://registry.npmjs.org/@redis/json/-/json-1.0.7.tgz",
"integrity": "sha512-6UyXfjVaTBTJtKNG4/9Z8PSpKE6XgSyEb8iwaqDcy+uKrd/DGYHTWkUdnQDyzm727V7p21WUMhsqz5oy65kPcQ==",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/search": {
"version": "1.2.0",
"resolved": "https://registry.npmjs.org/@redis/search/-/search-1.2.0.tgz",
"integrity": "sha512-tYoDBbtqOVigEDMAcTGsRlMycIIjwMCgD8eR2t0NANeQmgK/lvxNAvYyb6bZDD4frHRhIHkJu2TBRvB0ERkOmw==",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@redis/time-series": {
"version": "1.1.0",
"resolved": "https://registry.npmjs.org/@redis/time-series/-/time-series-1.1.0.tgz",
"integrity": "sha512-c1Q99M5ljsIuc4YdaCwfUEXsofakb9c8+Zse2qxTadu8TalLXuAESzLvFAvNVbkmSlvlzIQOLpBCmWI9wTOt+g==",
"peerDependencies": {
"@redis/client": "^1.0.0"
}
},
"node_modules/@rushstack/eslint-patch": {
"version": "1.10.3",
"resolved": "https://registry.npmjs.org/@rushstack/eslint-patch/-/eslint-patch-1.10.3.tgz",
@ -6715,6 +6775,16 @@
"@types/react": "*"
}
},
"node_modules/@types/redis": {
"version": "4.0.11",
"resolved": "https://registry.npmjs.org/@types/redis/-/redis-4.0.11.tgz",
"integrity": "sha512-bI+gth8La8Wg/QCR1+V1fhrL9+LZUSWfcqpOj2Kc80ZQ4ffbdL173vQd5wovmoV9i071FU9oP2g6etLuEwb6Rg==",
"deprecated": "This is a stub types definition. redis provides its own type definitions, so you do not need this installed.",
"dev": true,
"dependencies": {
"redis": "*"
}
},
"node_modules/@types/retry": {
"version": "0.12.0",
"resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.0.tgz",
@ -7832,6 +7902,14 @@
"node": ">=6"
}
},
"node_modules/cluster-key-slot": {
"version": "1.1.2",
"resolved": "https://registry.npmjs.org/cluster-key-slot/-/cluster-key-slot-1.1.2.tgz",
"integrity": "sha512-RMr0FhtfXemyinomL4hrWcYJxmX6deFdCxpJzhDttxgO1+bcCnkk+9drydLVDmAMG7NE6aN/fl4F7ucU/90gAA==",
"engines": {
"node": ">=0.10.0"
}
},
"node_modules/code-red": {
"version": "1.0.4",
"resolved": "https://registry.npmjs.org/code-red/-/code-red-1.0.4.tgz",
@ -9416,6 +9494,14 @@
"resolved": "https://registry.npmjs.org/fzy.js/-/fzy.js-0.4.1.tgz",
"integrity": "sha512-4sPVXf+9oGhzg2tYzgWe4hgAY0wEbkqeuKVEgdnqX8S8VcLosQsDjb0jV+f5uoQlf8INWId1w0IGoufAoik1TA=="
},
"node_modules/generic-pool": {
"version": "3.9.0",
"resolved": "https://registry.npmjs.org/generic-pool/-/generic-pool-3.9.0.tgz",
"integrity": "sha512-hymDOu5B53XvN4QT9dBmZxPX4CWhBPPLguTZ9MMFeFa/Kg0xWVfylOVNlJji/E7yTZWFd/q9GO5TxDLq156D7g==",
"engines": {
"node": ">= 4"
}
},
"node_modules/gensync": {
"version": "1.0.0-beta.2",
"resolved": "https://registry.npmjs.org/gensync/-/gensync-1.0.0-beta.2.tgz",
@ -12929,6 +13015,22 @@
"node": ">=8.10.0"
}
},
"node_modules/redis": {
"version": "4.7.0",
"resolved": "https://registry.npmjs.org/redis/-/redis-4.7.0.tgz",
"integrity": "sha512-zvmkHEAdGMn+hMRXuMBtu4Vo5P6rHQjLoHftu+lBqq8ZTA3RCVC/WzD790bkKKiNFp7d5/9PcSD19fJyyRvOdQ==",
"workspaces": [
"./packages/*"
],
"dependencies": {
"@redis/bloom": "1.2.0",
"@redis/client": "1.6.0",
"@redis/graph": "1.1.1",
"@redis/json": "1.0.7",
"@redis/search": "1.2.0",
"@redis/time-series": "1.1.0"
}
},
"node_modules/reflect.getprototypeof": {
"version": "1.0.6",
"resolved": "https://registry.npmjs.org/reflect.getprototypeof/-/reflect.getprototypeof-1.0.6.tgz",

View file

@ -37,6 +37,7 @@
"react-dom": "^18.3.1",
"react-markdown": "^9.0.1",
"react-resizable-panels": "^2.1.7",
"redis": "^4.7.0",
"remark-gfm": "^4.0.0",
"rowboat-shared": "github:rowboatlabs/shared",
"sharp": "^0.33.4",
@ -53,6 +54,7 @@
"@types/node": "^20",
"@types/react": "^18",
"@types/react-dom": "^18",
"@types/redis": "^4.0.11",
"eslint": "^8",
"eslint-config-next": "14.2.5",
"postcss": "^8",

View file

@ -23,6 +23,9 @@ services:
- AUTH0_ISSUER_BASE_URL=${AUTH0_ISSUER_BASE_URL}
- AUTH0_CLIENT_ID=${AUTH0_CLIENT_ID}
- AUTH0_CLIENT_SECRET=${AUTH0_CLIENT_SECRET}
- REDIS_URL=redis://redis:6379
- MAX_QUERIES_PER_MINUTE=${MAX_QUERIES_PER_MINUTE}
- MAX_PROJECTS_PER_USER=${MAX_PROJECTS_PER_USER}
restart: unless-stopped
agents:
@ -54,3 +57,9 @@ services:
ports:
- "8000:8000"
restart: unless-stopped
redis:
image: redis:latest
ports:
- "6379:6379"
restart: unless-stopped