Bugfix/Rate limit updating (#3311)

fix rate limit updating
FlowiseAI · Oct 5, 2024 · 10bfba7 · 10bfba7
1 parent f5cedb2
commit 10bfba7
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 22 deletions.
diff --git 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
@@ -21,28 +21,35 @@ 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()
     const idRateLimiter = rateLimiters[id]
     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
@@ -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 = () => {
                     }
                 />
             </Typography>
-            {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')}
-
-            <StyledButton style={{ marginBottom: 10, marginTop: 10 }} variant='contained' onClick={() => onSave()}>
+            <SwitchInput label='Enable Rate Limit' onChange={handleChange} value={rateLimitStatus} />
+            {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')}
+                </>
+            )}
+            <StyledButton
+                disabled={checkDisabled()}
+                style={{ marginBottom: 10, marginTop: 10 }}
+                variant='contained'
+                onClick={() => onSave()}
+            >
                 Save Changes
             </StyledButton>
         </>