diff options
| -rw-r--r-- | packages/console/app/src/routes/zen/util/rateLimiter.ts | 50 | ||||
| -rw-r--r-- | packages/console/core/src/subscription.ts | 1 |
2 files changed, 33 insertions, 18 deletions
diff --git a/packages/console/app/src/routes/zen/util/rateLimiter.ts b/packages/console/app/src/routes/zen/util/rateLimiter.ts index 152ed9b8d..042ec51f2 100644 --- a/packages/console/app/src/routes/zen/util/rateLimiter.ts +++ b/packages/console/app/src/routes/zen/util/rateLimiter.ts @@ -17,36 +17,52 @@ export function createRateLimiter( const dict = i18n(localeFromRequest(request)) const limits = Subscription.getFreeLimits() - const limitValue = - limits.checkHeader && !request.headers.get(limits.checkHeader) - ? limits.fallbackValue - : (rateLimit ?? limits.dailyRequests) + const headerExists = request.headers.has(limits.checkHeader) + const dailyLimit = !headerExists ? limits.fallbackValue : (rateLimit ?? limits.dailyRequests) + const isDefaultModel = headerExists && !rateLimit const ip = !rawIp.length ? "unknown" : rawIp const now = Date.now() - const interval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) + const lifetimeInterval = "" + const dailyInterval = rateLimit ? `${buildYYYYMMDD(now)}${modelId.substring(0, 2)}` : buildYYYYMMDD(now) + + let _isNew: boolean return { - track: async () => { - await Database.use((tx) => - tx - .insert(IpRateLimitTable) - .values({ ip, interval, count: 1 }) - .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), - ) - }, check: async () => { const rows = await Database.use((tx) => tx .select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count }) .from(IpRateLimitTable) - .where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, [interval]))), + .where( + and( + eq(IpRateLimitTable.ip, ip), + isDefaultModel + ? inArray(IpRateLimitTable.interval, [lifetimeInterval, dailyInterval]) + : inArray(IpRateLimitTable.interval, [dailyInterval]), + ), + ), ) - const total = rows.reduce((sum, r) => sum + r.count, 0) - logger.debug(`rate limit total: ${total}`) - if (total >= limitValue) + const lifetimeCount = rows.find((r) => r.interval === lifetimeInterval)?.count ?? 0 + const dailyCount = rows.find((r) => r.interval === dailyInterval)?.count ?? 0 + logger.debug(`rate limit lifetime: ${lifetimeCount}, daily: ${dailyCount}`) + + _isNew = isDefaultModel && lifetimeCount < dailyLimit * 7 + + if ((_isNew && dailyCount >= dailyLimit * 2) || (!_isNew && dailyCount >= dailyLimit)) throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now)) }, + track: async () => { + await Database.use((tx) => + tx + .insert(IpRateLimitTable) + .values([ + { ip, interval: dailyInterval, count: 1 }, + ...(_isNew ? [{ ip, interval: lifetimeInterval, count: 1 }] : []), + ]) + .onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }), + ) + }, } } diff --git a/packages/console/core/src/subscription.ts b/packages/console/core/src/subscription.ts index 994feaff0..9d6c3ce2b 100644 --- a/packages/console/core/src/subscription.ts +++ b/packages/console/core/src/subscription.ts @@ -8,7 +8,6 @@ export namespace Subscription { const LimitsSchema = z.object({ free: z.object({ promoTokens: z.number().int(), - newDailyRequests: z.number().int(), dailyRequests: z.number().int(), checkHeader: z.string(), fallbackValue: z.number().int(), |