Merge pull request #843 from FlowiseAI/feature/RateLimit

feature/RateLimit
2023-09-10 21:04:14 +08:00 · 2023-09-10 21:04:14 +08:00 · 0a5d453338
parent 3a813fa044 c4243d6bd1
commit 0a5d453338
16 changed files with 299 additions and 11 deletions
--- a/docker/.env.example
+++ b/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=""
--- a/packages/server/.env.example
+++ b/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=""
--- a/packages/server/package.json
+++ b/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",
--- a/packages/server/src/Interface.ts
+++ b/packages/server/src/Interface.ts
@ -15,6 +15,7 @@ export interface IChatFlow {
    isPublic?: boolean
    apikeyid?: string
    chatbotConfig?: string
+    apiConfig?: any
 }

 export interface IChatMessage {
--- a/packages/server/src/database/entities/ChatFlow.ts
+++ b/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

--- a/packages/server/src/database/migrations/mysql/1694099200729-AddApiConfig.ts
+++ b/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\`;`)
+    }
+}
--- a/packages/server/src/database/migrations/mysql/index.ts
+++ b/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
 ]
--- a/packages/server/src/database/migrations/postgres/1694099183389-AddApiConfig.ts
+++ b/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";`)
+    }
+}
--- a/packages/server/src/database/migrations/postgres/index.ts
+++ b/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
 ]
--- a/packages/server/src/database/migrations/sqlite/1694090982460-AddApiConfig.ts
+++ b/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";`)
+    }
+}
--- a/packages/server/src/database/migrations/sqlite/index.ts
+++ b/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
 ]
--- a/packages/server/src/index.ts
+++ b/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) => {
-            await this.processPrediction(req, res, socketIO)
-        })
+        this.app.post(
+            '/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()

--- a/packages/server/src/utils/rateLimit.ts
+++ b/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)
+    })
+}
--- a/packages/ui/src/assets/images/settings.svg
+++ b/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>
--- a/packages/ui/src/views/chatflows/APICodeDialog.js
+++ b/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>
--- a/packages/ui/src/views/chatflows/Configuration.js
+++ b/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