Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bugfix/Rate limit updating #3311

Merged
merged 1 commit into from
Oct 5, 2024
Merged

Bugfix/Rate limit updating #3311

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions packages/server/src/controllers/chatflows/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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)
Expand Down
17 changes: 12 additions & 5 deletions packages/server/src/utils/rateLimit.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
)
}
58 changes: 43 additions & 15 deletions packages/ui/src/ui-component/extended/RateLimit.jsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -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, {
Expand Down Expand Up @@ -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>
</>
Expand Down
Loading