From 0da6cd9a99d56e64a2914717c44dad23a209df91 Mon Sep 17 00:00:00 2001
From: Ronny Gunawan <3048897+ronnygunawan@users.noreply.github.com>
Date: Sun, 3 Dec 2023 19:31:59 +0700
Subject: [PATCH] Stricter rate limiting

---
 BotNet.Services/BotCommands/OpenAI.cs | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/BotNet.Services/BotCommands/OpenAI.cs b/BotNet.Services/BotCommands/OpenAI.cs
index f0b0fc1..a427cf0 100644
--- a/BotNet.Services/BotCommands/OpenAI.cs
+++ b/BotNet.Services/BotCommands/OpenAI.cs
@@ -4,6 +4,7 @@
 using System.IO;
 using System.Linq;
 using System.Net;
+using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 using BotNet.Services.MarkdownV2;
@@ -765,7 +766,8 @@ await botClient.SendTextMessageAsync(
 			}
 		}
 
-		private static readonly RateLimiter IMAGE_GENERATION_RATE_LIMITER = RateLimiter.PerUser(1, TimeSpan.FromMinutes(3));
+		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));
 		public static async Task StreamChatWithFriendlyBotAsync(
 			ITelegramBotClient botClient,
 			IServiceProvider serviceProvider,
@@ -829,7 +831,11 @@ await serviceProvider.GetRequiredService<FriendlyBot>().StreamChatAsync(
 							);
 							break;
 						case ChatIntent.ImageGeneration:
-							IMAGE_GENERATION_RATE_LIMITER.ValidateActionRate(
+							IMAGE_GENERATION_PER_USER_RATE_LIMITER.ValidateActionRate(
+								chatId: message.Chat.Id,
+								userId: message.From.Id
+							);
+							IMAGE_GENERATION_PER_CHAT_RATE_LIMITER.ValidateActionRate(
 								chatId: message.Chat.Id,
 								userId: message.From.Id
 							);
@@ -881,6 +887,13 @@ await botClient.SendTextMessageAsync(
 						),
 						cancellationToken: cancellationToken);
 				}
+			} catch (HttpRequestException exc) when (exc.StatusCode == HttpStatusCode.TooManyRequests) {
+				await botClient.SendTextMessageAsync(
+					chatId: message.Chat.Id,
+					text: "<code>Too many requests.</code>",
+					parseMode: ParseMode.Html,
+					replyToMessageId: message.MessageId,
+					cancellationToken: cancellationToken);
 			} catch (OperationCanceledException) {
 				await botClient.SendTextMessageAsync(
 					chatId: message.Chat.Id,