Merge pull request #843 from FlowiseAI/feature/RateLimit

feature/RateLimit

docker/.env.example

@@ -5,6 +5,8 @@ APIKEY_PATH=/root/.flowise
 SECRETKEY_PATH=/root/.flowise
 LOG_PATH=/root/.flowise/logs
 
+# NUMBER_OF_PROXIES= 1
+
 # DATABASE_TYPE=postgres
 # DATABASE_PORT=""
 # DATABASE_HOST=""

packages/server/.env.example

@@ -5,6 +5,8 @@ PASSPHRASE=MYPASSPHRASE # Passphrase used to create encryption key
 # SECRETKEY_PATH=/your_api_key_path/.flowise
 # LOG_PATH=/your_log_path/.flowise/logs
 
+# NUMBER_OF_PROXIES= 1
+
 # DATABASE_TYPE=postgres
 # DATABASE_PORT=""
 # DATABASE_HOST=""

packages/server/package.json

@@ -46,12 +46,14 @@
     "license": "SEE LICENSE IN LICENSE.md",
     "dependencies": {
         "@oclif/core": "^1.13.10",
+        "async-mutex": "^0.4.0",
         "axios": "^0.27.2",
         "cors": "^2.8.5",
         "crypto-js": "^4.1.1",
         "dotenv": "^16.0.0",
         "express": "^4.17.3",
         "express-basic-auth": "^1.2.1",
+        "express-rate-limit": "^6.9.0",
         "flowise-components": "*",
         "flowise-ui": "*",
         "moment-timezone": "^0.5.34",

packages/server/src/Interface.ts

@@ -15,6 +15,7 @@ export interface IChatFlow {
     isPublic?: boolean
     apikeyid?: string
     chatbotConfig?: string
+    apiConfig?: any
 }
 
 export interface IChatMessage {

packages/server/src/database/entities/ChatFlow.ts

@@ -25,6 +25,9 @@ export class ChatFlow implements IChatFlow {
     @Column({ nullable: true, type: 'text' })
     chatbotConfig?: string
 
+    @Column({ nullable: true, type: 'text' })
+    apiConfig?: string
+
     @CreateDateColumn()
     createdDate: Date
 

packages/server/src/database/migrations/mysql/1694099200729-AddApiConfig.ts

@@ -0,0 +1,11 @@
+import { MigrationInterface, QueryRunner } from 'typeorm'
+
+export class AddApiConfig1694099200729 implements MigrationInterface {
+    public async up(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`ALTER TABLE \`chat_flow\` ADD COLUMN \`apiConfig\` TEXT;`)
+    }
+
+    public async down(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`ALTER TABLE \`chat_flow\` DROP COLUMN \`apiConfig\`;`)
+    }
+}

packages/server/src/database/migrations/mysql/index.ts

@@ -3,11 +3,13 @@ import { ModifyChatFlow1693997791471 } from './1693997791471-ModifyChatFlow'
 import { ModifyChatMessage1693999022236 } from './1693999022236-ModifyChatMessage'
 import { ModifyCredential1693999261583 } from './1693999261583-ModifyCredential'
 import { ModifyTool1694001465232 } from './1694001465232-ModifyTool'
+import { AddApiConfig1694099200729 } from './1694099200729-AddApiConfig'
 
 export const mysqlMigrations = [
     Init1693840429259,
     ModifyChatFlow1693997791471,
     ModifyChatMessage1693999022236,
     ModifyCredential1693999261583,
-    ModifyTool1694001465232
+    ModifyTool1694001465232,
+    AddApiConfig1694099200729
 ]

packages/server/src/database/migrations/postgres/1694099183389-AddApiConfig.ts

@@ -0,0 +1,11 @@
+import { MigrationInterface, QueryRunner } from 'typeorm'
+
+export class AddApiConfig1694099183389 implements MigrationInterface {
+    public async up(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`ALTER TABLE "chat_flow" ADD COLUMN "apiConfig" TEXT;`)
+    }
+
+    public async down(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`ALTER TABLE "chat_flow" DROP COLUMN "apiConfig";`)
+    }
+}

packages/server/src/database/migrations/postgres/index.ts

@@ -3,11 +3,13 @@ import { ModifyChatFlow1693995626941 } from './1693995626941-ModifyChatFlow'
 import { ModifyChatMessage1693996694528 } from './1693996694528-ModifyChatMessage'
 import { ModifyCredential1693997070000 } from './1693997070000-ModifyCredential'
 import { ModifyTool1693997339912 } from './1693997339912-ModifyTool'
+import { AddApiConfig1694099183389 } from './1694099183389-AddApiConfig'
 
 export const postgresMigrations = [
     Init1693891895163,
     ModifyChatFlow1693995626941,
     ModifyChatMessage1693996694528,
     ModifyCredential1693997070000,
-    ModifyTool1693997339912
+    ModifyTool1693997339912,
+    AddApiConfig1694099183389
 ]

packages/server/src/database/migrations/sqlite/1694090982460-AddApiConfig.ts

@@ -0,0 +1,11 @@
+import { MigrationInterface, QueryRunner } from 'typeorm'
+
+export class AddApiConfig1694090982460 implements MigrationInterface {
+    public async up(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`ALTER TABLE "chat_flow" ADD COLUMN "apiConfig" TEXT;`)
+    }
+
+    public async down(queryRunner: QueryRunner): Promise<void> {
+        await queryRunner.query(`ALTER TABLE "chat_flow" DROP COLUMN "apiConfig";`)
+    }
+}

packages/server/src/database/migrations/sqlite/index.ts

@@ -3,11 +3,13 @@ import { ModifyChatFlow1693920824108 } from './1693920824108-ModifyChatFlow'
 import { ModifyChatMessage1693921865247 } from './1693921865247-ModifyChatMessage'
 import { ModifyCredential1693923551694 } from './1693923551694-ModifyCredential'
 import { ModifyTool1693924207475 } from './1693924207475-ModifyTool'
+import { AddApiConfig1694090982460 } from './1694090982460-AddApiConfig'
 
 export const sqliteMigrations = [
     Init1693835579790,
     ModifyChatFlow1693920824108,
     ModifyChatMessage1693921865247,
     ModifyCredential1693923551694,
-    ModifyTool1693924207475
+    ModifyTool1693924207475,
+    AddApiConfig1694090982460
 ]

packages/server/src/index.ts

@@ -1,4 +1,4 @@
-import express, { Request, Response } from 'express'
+import express, { NextFunction, Request, Response } from 'express'
 import multer from 'multer'
 import path from 'path'
 import cors from 'cors'
@@ -54,6 +54,7 @@ import { Credential } from './database/entities/Credential'
 import { Tool } from './database/entities/Tool'
 import { ChatflowPool } from './ChatflowPool'
 import { ICommonObject, INodeOptionsValue } from 'flowise-components'
+import { createRateLimiter, getRateLimiter, initializeRateLimiter } from './utils/rateLimit'
 
 export class App {
     app: express.Application
@@ -86,6 +87,10 @@ export class App {
 
                 // Initialize encryption key
                 await getEncryptionKey()
+
+                // Initialize Rate Limit
+                const AllChatFlow: IChatFlow[] = await getAllChatFlow()
+                await initializeRateLimiter(AllChatFlow)
             })
             .catch((err) => {
                 logger.error('❌ [server]: Error during Data Source initialization:', err)
@@ -97,6 +102,9 @@ export class App {
         this.app.use(express.json({ limit: '50mb' }))
         this.app.use(express.urlencoded({ limit: '50mb', extended: true }))
 
+        if (process.env.NUMBER_OF_PROXIES && parseInt(process.env.NUMBER_OF_PROXIES) > 0)
+            this.app.set('trust proxy', parseInt(process.env.NUMBER_OF_PROXIES))
+
         // Allow access from *
         this.app.use(cors())
 
@@ -116,7 +124,8 @@ export class App {
                 '/api/v1/prediction/',
                 '/api/v1/node-icon/',
                 '/api/v1/components-credentials-icon/',
-                '/api/v1/chatflows-streaming'
+                '/api/v1/chatflows-streaming',
+                '/api/v1/ip'
             ]
             this.app.use((req, res, next) => {
                 if (req.url.includes('/api/v1/')) {
@@ -127,6 +136,16 @@ export class App {
 
         const upload = multer({ dest: `${path.join(__dirname, '..', 'uploads')}/` })
 
+        // ----------------------------------------
+        // Configure number of proxies in Host Environment
+        // ----------------------------------------
+        this.app.get('/api/v1/ip', (request, response) => {
+            response.send({
+                ip: request.ip,
+                msg: 'See the IP returned in the response. If it matches your IP address (which you can get by going to http://ip.nfriedly.com/ or https://api.ipify.org/), then the number of proxies is correct and the rate limiter should now work correctly. If not, then keep increasing the number until it does.'
+            })
+        })
+
         // ----------------------------------------
         // Components
         // ----------------------------------------
@@ -248,7 +267,7 @@ export class App {
 
         // Get all chatflows
         this.app.get('/api/v1/chatflows', async (req: Request, res: Response) => {
-            const chatflows: IChatFlow[] = await this.AppDataSource.getRepository(ChatFlow).find()
+            const chatflows: IChatFlow[] = await getAllChatFlow()
             return res.json(chatflows)
         })
 
@@ -317,6 +336,9 @@ export class App {
             const updateChatFlow = new ChatFlow()
             Object.assign(updateChatFlow, body)
 
+            updateChatFlow.id = chatflow.id
+            createRateLimiter(updateChatFlow)
+
             this.AppDataSource.getRepository(ChatFlow).merge(chatflow, updateChatFlow)
             const result = await this.AppDataSource.getRepository(ChatFlow).save(chatflow)
 
@@ -658,9 +680,14 @@ export class App {
         // ----------------------------------------
 
         // Send input message and get prediction result (External)
-        this.app.post('/api/v1/prediction/:id', upload.array('files'), async (req: Request, res: Response) => {
+        this.app.post(
-            await this.processPrediction(req, res, socketIO)
+            '/api/v1/prediction/:id',
-        })
+            upload.array('files'),
+            (req: Request, res: Response, next: NextFunction) => getRateLimiter(req, res, next),
+            async (req: Request, res: Response) => {
+                await this.processPrediction(req, res, socketIO)
+            }
+        )
 
         // Send input message and get prediction result (Internal)
         this.app.post('/api/v1/internal-prediction/:id', async (req: Request, res: Response) => {
@@ -999,6 +1026,10 @@ export async function getChatId(chatflowid: string) {
 
 let serverApp: App | undefined
 
+export async function getAllChatFlow(): Promise<IChatFlow[]> {
+    return await getDataSource().getRepository(ChatFlow).find()
+}
+
 export async function start(): Promise<void> {
     serverApp = new App()
 

packages/server/src/utils/rateLimit.ts

@@ -0,0 +1,47 @@
+import { NextFunction, Request, Response } from 'express'
+import { rateLimit, RateLimitRequestHandler } from 'express-rate-limit'
+import { IChatFlow } from '../Interface'
+import { Mutex } from 'async-mutex'
+
+let rateLimiters: Record<string, RateLimitRequestHandler> = {}
+const rateLimiterMutex = new Mutex()
+
+async function addRateLimiter(id: string, duration: number, limit: number, message: string) {
+    const release = await rateLimiterMutex.acquire()
+    try {
+        rateLimiters[id] = rateLimit({
+            windowMs: duration * 1000,
+            max: limit,
+            handler: (req, res) => {
+                res.status(429).send(message)
+            }
+        })
+    } finally {
+        release()
+    }
+}
+
+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) {
+    if (!chatFlow.apiConfig) return
+    const apiConfig: any = JSON.parse(chatFlow.apiConfig)
+    const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string } = apiConfig.rateLimit
+    if (!rateLimit) return
+    const { limitDuration, limitMax, limitMsg } = rateLimit
+    if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg)
+}
+
+export async function initializeRateLimiter(chatFlowPool: IChatFlow[]) {
+    await chatFlowPool.map(async (chatFlow) => {
+        await createRateLimiter(chatFlow)
+    })
+}

packages/ui/src/assets/images/settings.svg

@@ -0,0 +1,5 @@
+<svg xmlns="http://www.w3.org/2000/svg" class="icon icon-tabler icon-tabler-settings" width="24" height="24" viewBox="0 0 24 24" stroke-width="2" stroke="currentColor" fill="none" stroke-linecap="round" stroke-linejoin="round">
+   <path stroke="none" d="M0 0h24v24H0z" fill="none"></path>
+   <path d="M10.325 4.317c.426 -1.756 2.924 -1.756 3.35 0a1.724 1.724 0 0 0 2.573 1.066c1.543 -.94 3.31 .826 2.37 2.37a1.724 1.724 0 0 0 1.065 2.572c1.756 .426 1.756 2.924 0 3.35a1.724 1.724 0 0 0 -1.066 2.573c.94 1.543 -.826 3.31 -2.37 2.37a1.724 1.724 0 0 0 -2.572 1.065c-.426 1.756 -2.924 1.756 -3.35 0a1.724 1.724 0 0 0 -2.573 -1.066c-1.543 .94 -3.31 -.826 -2.37 -2.37a1.724 1.724 0 0 0 -1.065 -2.572c-1.756 -.426 -1.756 -2.924 0 -3.35a1.724 1.724 0 0 0 1.066 -2.573c-.94 -1.543 .826 -3.31 2.37 -2.37c1 .608 2.296 .07 2.572 -1.065z"></path>
+   <path d="M9 12a3 3 0 1 0 6 0a3 3 0 0 0 -6 0"></path>
+</svg>

packages/ui/src/views/chatflows/APICodeDialog.js

@@ -34,6 +34,7 @@ import javascriptSVG from 'assets/images/javascript.svg'
 import cURLSVG from 'assets/images/cURL.svg'
 import EmbedSVG from 'assets/images/embed.svg'
 import ShareChatbotSVG from 'assets/images/sharing.png'
+import settingsSVG from 'assets/images/settings.svg'
 
 // API
 import apiKeyApi from 'api/apikey'
@@ -46,6 +47,7 @@ import { CheckboxInput } from 'ui-component/checkbox/Checkbox'
 import { TableViewOnly } from 'ui-component/table/Table'
 
 import { IconBulb } from '@tabler/icons'
+import Configuration from './Configuration'
 
 function TabPanel(props) {
     const { children, value, index, ...other } = props
@@ -141,7 +143,7 @@ const APICodeDialog = ({ show, dialogProps, onCancel }) => {
     const navigate = useNavigate()
     const dispatch = useDispatch()
 
-    const codes = ['Embed', 'Python', 'JavaScript', 'cURL', 'Share Chatbot']
+    const codes = ['Embed', 'Python', 'JavaScript', 'cURL', 'Share Chatbot', 'Configuration']
     const [value, setValue] = useState(0)
     const [keyOptions, setKeyOptions] = useState([])
     const [apiKeys, setAPIKeys] = useState([])
@@ -321,6 +323,8 @@ query({"question": "Hey, how are you?"}).then((response) => {
             return cURLSVG
         } else if (codeLang === 'Share Chatbot') {
             return ShareChatbotSVG
+        } else if (codeLang === 'Configuration') {
+            return settingsSVG
         }
         return pythonSVG
     }
@@ -647,7 +651,7 @@ formData.append("openAIApiKey[openAIEmbeddings_0]", "sk-my-openai-2nd-key")`
                             </>
                         )}
                         {codeLang === 'Embed' && !chatflowApiKeyId && <EmbedChat chatflowid={dialogProps.chatflowid} />}
-                        {codeLang !== 'Embed' && codeLang !== 'Share Chatbot' && (
+                        {codeLang !== 'Embed' && codeLang !== 'Share Chatbot' && codeLang !== 'Configuration' && (
                             <>
                                 <CopyBlock
                                     theme={atomOneDark}
@@ -770,6 +774,7 @@ formData.append("openAIApiKey[openAIEmbeddings_0]", "sk-my-openai-2nd-key")`
                         {codeLang === 'Share Chatbot' && !chatflowApiKeyId && (
                             <ShareChatbot isSessionMemory={dialogProps.isSessionMemory} />
                         )}
+                        {codeLang === 'Configuration' && <Configuration />}
                     </TabPanel>
                 ))}
             </DialogContent>

packages/ui/src/views/chatflows/Configuration.js

@@ -0,0 +1,151 @@
+import { useState } from 'react'
+import { useDispatch, useSelector } from 'react-redux'
+import { enqueueSnackbar as enqueueSnackbarAction, closeSnackbar as closeSnackbarAction, SET_CHATFLOW } from 'store/actions'
+import PropTypes from 'prop-types'
+
+import { Box, Typography, Button, OutlinedInput } from '@mui/material'
+
+// Project import
+import { StyledButton } from 'ui-component/button/StyledButton'
+
+// Icons
+import { IconX } from '@tabler/icons'
+
+// API
+import chatflowsApi from 'api/chatflows'
+
+// utils
+import useNotifier from 'utils/useNotifier'
+
+const Configuration = () => {
+    const dispatch = useDispatch()
+    const chatflow = useSelector((state) => state.canvas.chatflow)
+    const chatflowid = chatflow.id
+    const apiConfig = chatflow.apiConfig ? JSON.parse(chatflow.apiConfig) : {}
+
+    useNotifier()
+
+    const enqueueSnackbar = (...args) => dispatch(enqueueSnackbarAction(...args))
+    const closeSnackbar = (...args) => dispatch(closeSnackbarAction(...args))
+
+    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: {}
+        }
+        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
+            }
+        }
+
+        return obj
+    }
+
+    const onSave = async () => {
+        try {
+            const saveResp = await chatflowsApi.updateChatflow(chatflowid, {
+                apiConfig: JSON.stringify(formatObj())
+            })
+            if (saveResp.data) {
+                enqueueSnackbar({
+                    message: 'API Configuration Saved',
+                    options: {
+                        key: new Date().getTime() + Math.random(),
+                        variant: 'success',
+                        action: (key) => (
+                            <Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
+                                <IconX />
+                            </Button>
+                        )
+                    }
+                })
+                dispatch({ type: SET_CHATFLOW, chatflow: saveResp.data })
+            }
+        } catch (error) {
+            console.error(error)
+            const errorData = error.response
+                ? error.response.data || `${error.response.status}: ${error.response.statusText}`
+                : error.message
+            enqueueSnackbar({
+                message: `Failed to save API Configuration: ${errorData}`,
+                options: {
+                    key: new Date().getTime() + Math.random(),
+                    variant: 'error',
+                    persist: true,
+                    action: (key) => (
+                        <Button style={{ color: 'white' }} onClick={() => closeSnackbar(key)}>
+                            <IconX />
+                        </Button>
+                    )
+                }
+            })
+        }
+    }
+
+    const onTextChanged = (value, fieldName) => {
+        switch (fieldName) {
+            case 'limitMax':
+                setLimitMax(value)
+                break
+            case 'limitDuration':
+                setLimitDuration(value)
+                break
+            case 'limitMsg':
+                setLimitMsg(value)
+                break
+        }
+    }
+
+    const textField = (message, fieldName, fieldLabel, fieldType = 'string', placeholder = '') => {
+        return (
+            <Box sx={{ pt: 2, pb: 2 }}>
+                <div style={{ display: 'flex', flexDirection: 'column', alignItems: 'flex-start' }}>
+                    <Typography sx={{ mb: 1 }}>{fieldLabel}</Typography>
+                    <OutlinedInput
+                        id={fieldName}
+                        type={fieldType}
+                        fullWidth
+                        value={message}
+                        placeholder={placeholder}
+                        name={fieldName}
+                        onChange={(e) => {
+                            onTextChanged(e.target.value, fieldName)
+                        }}
+                    />
+                </div>
+            </Box>
+        )
+    }
+
+    return (
+        <>
+            {/*Rate Limit*/}
+            <Typography variant='h4' sx={{ mb: 1, mt: 2 }}>
+                Rate Limit
+            </Typography>
+            {textField(limitMax, 'limitMax', 'Message Limit per Duration', 'number')}
+            {textField(limitDuration, 'limitDuration', 'Duration in Second', 'number')}
+            {textField(limitMsg, 'limitMsg', 'Limit Message', 'string')}
+
+            <StyledButton style={{ marginBottom: 10, marginTop: 10 }} variant='contained' onClick={() => onSave()}>
+                Save Changes
+            </StyledButton>
+        </>
+    )
+}
+
+Configuration.propTypes = {
+    isSessionMemory: PropTypes.bool
+}
+
+export default Configuration