Merge pull request #77 from teknologi-umum/ai-overhaul

Add global rate limit for image generation
teknologi-umum · Dec 3, 2023 · a0c045e · a0c045e
2 parents 9f3b54d + 01f0776
commit a0c045e
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/BotNet.Services/BotCommands/OpenAI.cs b/BotNet.Services/BotCommands/OpenAI.cs
@@ -767,7 +767,8 @@ await botClient.SendTextMessageAsync(
 		}
 
 		private static readonly RateLimiter IMAGE_GENERATION_PER_USER_RATE_LIMITER = RateLimiter.PerUser(1, TimeSpan.FromMinutes(10));
-		private static readonly RateLimiter IMAGE_GENERATION_PER_CHAT_RATE_LIMITER = RateLimiter.PerUser(2, TimeSpan.FromMinutes(3));
+		private static readonly RateLimiter IMAGE_GENERATION_PER_CHAT_RATE_LIMITER = RateLimiter.PerChat(2, TimeSpan.FromMinutes(5));
+		private static readonly RateLimiter IMAGE_GENERATION_GLOBAL_RATE_LIMITER = RateLimiter.PerChat(1, TimeSpan.FromMinutes(1));
 		public static async Task StreamChatWithFriendlyBotAsync(
 			ITelegramBotClient botClient,
 			IServiceProvider serviceProvider,
@@ -839,6 +840,10 @@ await serviceProvider.GetRequiredService<FriendlyBot>().StreamChatAsync(
 								chatId: message.Chat.Id,
 								userId: message.From.Id
 							);
+							IMAGE_GENERATION_GLOBAL_RATE_LIMITER.ValidateActionRate(
+								chatId: 0,
+								userId: 0
+							);
 							Message busyMessage = await botClient.SendTextMessageAsync(
 								chatId: message.Chat.Id,
 								text: "Generating image… ⏳",