From f8fd49a00f7d8282b388b86138e4be1d4e1a446d Mon Sep 17 00:00:00 2001 From: Henry Date: Fri, 4 Oct 2024 17:34:29 +0100 Subject: [PATCH] fix rate limit updating --- .../server/src/controllers/chatflows/index.ts | 4 +- packages/server/src/utils/rateLimit.ts | 17 ++++-- .../src/ui-component/extended/RateLimit.jsx | 58 ++++++++++++++----- 3 files changed, 57 insertions(+), 22 deletions(-) diff --git a/packages/server/src/controllers/chatflows/index.ts b/packages/server/src/controllers/chatflows/index.ts index 523a6572739..0edaab40fb7 100644 --- a/packages/server/src/controllers/chatflows/index.ts +++ b/packages/server/src/controllers/chatflows/index.ts @@ -2,7 +2,7 @@ import { NextFunction, Request, Response } from 'express' import { StatusCodes } from 'http-status-codes' import apiKeyService from '../../services/apikey' import { ChatFlow } from '../../database/entities/ChatFlow' -import { createRateLimiter } from '../../utils/rateLimit' +import { updateRateLimiter } from '../../utils/rateLimit' import { InternalFlowiseError } from '../../errors/internalFlowiseError' import { ChatflowType } from '../../Interface' import chatflowsService from '../../services/chatflows' @@ -130,7 +130,7 @@ const updateChatflow = async (req: Request, res: Response, next: NextFunction) = Object.assign(updateChatFlow, body) updateChatFlow.id = chatflow.id - createRateLimiter(updateChatFlow) + updateRateLimiter(updateChatFlow) const apiResponse = await chatflowsService.updateChatflow(chatflow, updateChatFlow) return res.json(apiResponse) diff --git a/packages/server/src/utils/rateLimit.ts b/packages/server/src/utils/rateLimit.ts index 0c08f1d24e7..4cad150279e 100644 --- a/packages/server/src/utils/rateLimit.ts +++ b/packages/server/src/utils/rateLimit.ts @@ -21,6 +21,12 @@ async function addRateLimiter(id: string, duration: number, limit: number, messa } } +function removeRateLimit(id: string) { + if (rateLimiters[id]) { + delete rateLimiters[id] + } +} + export function getRateLimiter(req: Request, res: Response, next: NextFunction) { const id = req.params.id if (!rateLimiters[id]) return next() @@ -28,21 +34,22 @@ export function getRateLimiter(req: Request, res: Response, next: NextFunction) return idRateLimiter(req, res, next) } -export async function createRateLimiter(chatFlow: IChatFlow) { +export async function updateRateLimiter(chatFlow: IChatFlow) { if (!chatFlow.apiConfig) return const apiConfig = JSON.parse(chatFlow.apiConfig) - const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string } = apiConfig.rateLimit + const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string; status?: boolean } = apiConfig.rateLimit if (!rateLimit) return - const { limitDuration, limitMax, limitMsg } = rateLimit - if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg) + const { limitDuration, limitMax, limitMsg, status } = rateLimit + if (status === false) removeRateLimit(chatFlow.id) + else if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg) } export async function initializeRateLimiter(chatFlowPool: IChatFlow[]) { await Promise.all( chatFlowPool.map(async (chatFlow) => { - await createRateLimiter(chatFlow) + await updateRateLimiter(chatFlow) }) ) } diff --git a/packages/ui/src/ui-component/extended/RateLimit.jsx b/packages/ui/src/ui-component/extended/RateLimit.jsx index 99166d2144f..fef8063d73c 100644 --- a/packages/ui/src/ui-component/extended/RateLimit.jsx +++ b/packages/ui/src/ui-component/extended/RateLimit.jsx @@ -8,6 +8,7 @@ import { Box, Typography, Button, OutlinedInput } from '@mui/material' // Project import import { StyledButton } from '@/ui-component/button/StyledButton' import { TooltipWithParser } from '@/ui-component/tooltip/TooltipWithParser' +import { SwitchInput } from '@/ui-component/switch/Switch' // Icons import { IconX } from '@tabler/icons-react' @@ -29,29 +30,47 @@ const RateLimit = () => { const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args)) const closeSnackbar = (...args) => dispatch(closeSnackbarAction(...args)) + const [rateLimitStatus, setRateLimitStatus] = useState(false) const [limitMax, setLimitMax] = useState(apiConfig?.rateLimit?.limitMax ?? '') const [limitDuration, setLimitDuration] = useState(apiConfig?.rateLimit?.limitDuration ?? '') const [limitMsg, setLimitMsg] = useState(apiConfig?.rateLimit?.limitMsg ?? '') const formatObj = () => { const obj = { - rateLimit: {} + rateLimit: { status: rateLimitStatus } } - const rateLimitValuesBoolean = [!limitMax, !limitDuration, !limitMsg] - const rateLimitFilledValues = rateLimitValuesBoolean.filter((value) => value === false) - if (rateLimitFilledValues.length >= 1 && rateLimitFilledValues.length <= 2) { - throw new Error('Need to fill all rate limit input fields') - } else if (rateLimitFilledValues.length === 3) { - obj.rateLimit = { - limitMax, - limitDuration, - limitMsg + + if (rateLimitStatus) { + const rateLimitValuesBoolean = [!limitMax, !limitDuration, !limitMsg] + const rateLimitFilledValues = rateLimitValuesBoolean.filter((value) => value === false) + if (rateLimitFilledValues.length >= 1 && rateLimitFilledValues.length <= 2) { + throw new Error('Need to fill all rate limit input fields') + } else if (rateLimitFilledValues.length === 3) { + obj.rateLimit = { + ...obj.rateLimit, + limitMax, + limitDuration, + limitMsg + } } } return obj } + const handleChange = (value) => { + setRateLimitStatus(value) + } + + const checkDisabled = () => { + if (rateLimitStatus) { + if (limitMax === '' || limitDuration === '' || limitMsg === '') { + return true + } + } + return false + } + const onSave = async () => { try { const saveResp = await chatflowsApi.updateChatflow(chatflowid, { @@ -139,11 +158,20 @@ const RateLimit = () => { } /> - {textField(limitMax, 'limitMax', 'Message Limit per Duration', 'number', '5')} - {textField(limitDuration, 'limitDuration', 'Duration in Second', 'number', '60')} - {textField(limitMsg, 'limitMsg', 'Limit Message', 'string', 'You have reached the quota')} - - onSave()}> + + {rateLimitStatus && ( + <> + {textField(limitMax, 'limitMax', 'Message Limit per Duration', 'number', '5')} + {textField(limitDuration, 'limitDuration', 'Duration in Second', 'number', '60')} + {textField(limitMsg, 'limitMsg', 'Limit Message', 'string', 'You have reached the quota')} + + )} + onSave()} + > Save Changes