Limit image generation concurrency

teknologi-umum · Dec 3, 2023 · a16d81e · a16d81e
1 parent 458043b
commit a16d81e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 deletions.
diff --git a/BotNet.Services/BotCommands/OpenAI.cs b/BotNet.Services/BotCommands/OpenAI.cs
@@ -765,6 +765,7 @@ await botClient.SendTextMessageAsync(
 			}
 		}
 
+		private static readonly RateLimiter IMAGE_GENERATION_RATE_LIMITER = RateLimiter.PerUser(1, TimeSpan.FromMinutes(3));
 		public static async Task StreamChatWithFriendlyBotAsync(
 			ITelegramBotClient botClient,
 			IServiceProvider serviceProvider,
@@ -828,6 +829,10 @@ await serviceProvider.GetRequiredService<FriendlyBot>().StreamChatAsync(
 							);
 							break;
 						case ChatIntent.ImageGeneration:
+							IMAGE_GENERATION_RATE_LIMITER.ValidateActionRate(
+								chatId: message.Chat.Id,
+								userId: message.From.Id
+							);
 							Message busyMessage = await botClient.SendTextMessageAsync(
 								chatId: message.Chat.Id,
 								text: "Generating image… ⏳",

diff --git a/BotNet.Services/OpenAI/Skills/ImageGenerationBot.cs b/BotNet.Services/OpenAI/Skills/ImageGenerationBot.cs
@@ -7,14 +7,23 @@ public class ImageGenerationBot(
 		OpenAIClient openAIClient
 	) {
 		private readonly OpenAIClient _openAIClient = openAIClient;
+		private static readonly SemaphoreSlim SEMAPHORE = new(1, 1);
 
-		public Task<Uri> GenerateImageAsync(
+		public async Task<Uri> GenerateImageAsync(
 			string prompt,
 			CancellationToken cancellationToken
-		) => _openAIClient.GenerateImageAsync(
-			model: "dall-e-3",
-			prompt: prompt,
-			cancellationToken: cancellationToken
-		);
+		) {
+			// dall-e-3 endpoint does not allow concurrent requests
+			await SEMAPHORE.WaitAsync(cancellationToken);
+			try {
+				return await _openAIClient.GenerateImageAsync(
+					model: "dall-e-3",
+					prompt: prompt,
+					cancellationToken: cancellationToken
+				);
+			} finally {
+				SEMAPHORE.Release();
+			}
+		}
 	}
 }