Bugfix/Rate limit updating (#3311)

fix rate limit updating
This commit is contained in:
Henry Heng
2024-10-05 13:33:53 +01:00
committed by GitHub
parent f5cedb2460
commit 10bfba78dd
3 changed files with 57 additions and 22 deletions
@@ -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)
+12 -5
View File
@@ -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)
})
)
}