Merge pull request #957 from FlowiseAI/feature/RateLimit

Chore/RateLimit
This commit is contained in:
Henry Heng
2023-09-21 22:35:53 +01:00
committed by GitHub
2 changed files with 11 additions and 5 deletions
+2
View File
@@ -27,6 +27,7 @@ export default class Start extends Command {
LOG_LEVEL: Flags.string(), LOG_LEVEL: Flags.string(),
TOOL_FUNCTION_BUILTIN_DEP: Flags.string(), TOOL_FUNCTION_BUILTIN_DEP: Flags.string(),
TOOL_FUNCTION_EXTERNAL_DEP: Flags.string(), TOOL_FUNCTION_EXTERNAL_DEP: Flags.string(),
NUMBER_OF_PROXIES: Flags.string(),
DATABASE_TYPE: Flags.string(), DATABASE_TYPE: Flags.string(),
DATABASE_PATH: Flags.string(), DATABASE_PATH: Flags.string(),
DATABASE_PORT: Flags.string(), DATABASE_PORT: Flags.string(),
@@ -72,6 +73,7 @@ export default class Start extends Command {
if (flags.PORT) process.env.PORT = flags.PORT if (flags.PORT) process.env.PORT = flags.PORT
if (flags.DEBUG) process.env.DEBUG = flags.DEBUG if (flags.DEBUG) process.env.DEBUG = flags.DEBUG
if (flags.NUMBER_OF_PROXIES) process.env.NUMBER_OF_PROXIES = flags.NUMBER_OF_PROXIES
// Authorization // Authorization
if (flags.FLOWISE_USERNAME) process.env.FLOWISE_USERNAME = flags.FLOWISE_USERNAME if (flags.FLOWISE_USERNAME) process.env.FLOWISE_USERNAME = flags.FLOWISE_USERNAME
+7 -3
View File
@@ -12,7 +12,7 @@ async function addRateLimiter(id: string, duration: number, limit: number, messa
rateLimiters[id] = rateLimit({ rateLimiters[id] = rateLimit({
windowMs: duration * 1000, windowMs: duration * 1000,
max: limit, max: limit,
handler: (req, res) => { handler: (_, res) => {
res.status(429).send(message) res.status(429).send(message)
} }
}) })
@@ -33,15 +33,19 @@ export function getRateLimiter(req: Request, res: Response, next: NextFunction)
export async function createRateLimiter(chatFlow: IChatFlow) { export async function createRateLimiter(chatFlow: IChatFlow) {
if (!chatFlow.apiConfig) return if (!chatFlow.apiConfig) return
const apiConfig: any = JSON.parse(chatFlow.apiConfig) const apiConfig = JSON.parse(chatFlow.apiConfig)
const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string } = apiConfig.rateLimit const rateLimit: { limitDuration: number; limitMax: number; limitMsg: string } = apiConfig.rateLimit
if (!rateLimit) return if (!rateLimit) return
const { limitDuration, limitMax, limitMsg } = rateLimit const { limitDuration, limitMax, limitMsg } = rateLimit
if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg) if (limitMax && limitDuration && limitMsg) await addRateLimiter(chatFlow.id, limitDuration, limitMax, limitMsg)
} }
export async function initializeRateLimiter(chatFlowPool: IChatFlow[]) { export async function initializeRateLimiter(chatFlowPool: IChatFlow[]) {
await chatFlowPool.map(async (chatFlow) => { await Promise.all(
chatFlowPool.map(async (chatFlow) => {
await createRateLimiter(chatFlow) await createRateLimiter(chatFlow)
}) })
)
} }