Add authorization for image generation commands

teknologi-umum · Dec 31, 2023 · f603d8c · f603d8c
1 parent 0713c53
commit f603d8c
Show file tree

Hide file tree

Showing 7 changed files with 120 additions and 136 deletions.
diff --git a/BotNet.CommandHandlers/AI/OpenAI/OpenAIImageGenerationPromptHandler.cs b/BotNet.CommandHandlers/AI/OpenAI/OpenAIImageGenerationPromptHandler.cs
@@ -2,7 +2,6 @@
 using BotNet.Commands.AI.OpenAI;
 using BotNet.Commands.BotUpdate.Message;
 using BotNet.Services.OpenAI.Skills;
-using BotNet.Services.RateLimit;
 using Microsoft.Extensions.Logging;
 using Telegram.Bot;
 using Telegram.Bot.Types;
@@ -15,26 +14,12 @@ public sealed class OpenAIImageGenerationPromptHandler(
 		ITelegramMessageCache telegramMessageCache,
 		ILogger<OpenAIImageGenerationPromptHandler> logger
 	) : ICommandHandler<OpenAIImageGenerationPrompt> {
-		private static readonly RateLimiter IMAGE_GENERATION_RATE_LIMITER = RateLimiter.PerUser(1, TimeSpan.FromMinutes(5));
-
 		private readonly ITelegramBotClient _telegramBotClient = telegramBotClient;
 		private readonly ImageGenerationBot _imageGenerationBot = imageGenerationBot;
 		private readonly ITelegramMessageCache _telegramMessageCache = telegramMessageCache;
 		private readonly ILogger<OpenAIImageGenerationPromptHandler> _logger = logger;
 
 		public Task Handle(OpenAIImageGenerationPrompt command, CancellationToken cancellationToken) {
-			try {
-				IMAGE_GENERATION_RATE_LIMITER.ValidateActionRate(command.ChatId, command.SenderId);
-			} catch (RateLimitExceededException exc) {
-				return _telegramBotClient.SendTextMessageAsync(
-					chatId: command.ChatId,
-					text: $"Anda belum mendapat giliran. Coba lagi {exc.Cooldown}.",
-					parseMode: ParseMode.Html,
-					replyToMessageId: command.PromptMessageId,
-					cancellationToken: cancellationToken
-				);
-			}
-
 			// Fire and forget
 			Task.Run(async () => {
 				try {

diff --git a/BotNet.CommandHandlers/AI/OpenAI/OpenAITextPromptHandler.cs b/BotNet.CommandHandlers/AI/OpenAI/OpenAITextPromptHandler.cs
@@ -1,4 +1,5 @@
-using BotNet.Commands;
+using BotNet.CommandHandlers.Art;
+using BotNet.Commands;
 using BotNet.Commands.AI.OpenAI;
 using BotNet.Commands.AI.Stability;
 using BotNet.Commands.BotUpdate.Message;
@@ -81,6 +82,21 @@ select ChatMessage.FromText(
 
 				// Handle image generation intent
 				if (response.StartsWith("ImageGeneration:")) {
+					if (command.CommandPriority != CommandPriority.VIPChat) {
+						try {
+							ArtCommandHandler.IMAGE_GENERATION_RATE_LIMITER.ValidateActionRate(command.ChatId, command.SenderId);
+						} catch (RateLimitExceededException exc) {
+							await _telegramBotClient.SendTextMessageAsync(
+								chatId: command.ChatId,
+								text: $"Anda belum mendapat giliran. Coba lagi {exc.Cooldown}.",
+								parseMode: ParseMode.Html,
+								replyToMessageId: command.PromptMessageId,
+								cancellationToken: cancellationToken
+							);
+							return;
+						}
+					}
+
 					string imageGenerationPrompt = response.Substring(response.IndexOf(':') + 1).Trim();
 					switch (command.CommandPriority) {
 						case CommandPriority.VIPChat:

diff --git a/BotNet.CommandHandlers/AI/Stability/ArtCommandHandler.cs b/BotNet.CommandHandlers/AI/Stability/ArtCommandHandler.cs
diff --git a/BotNet.CommandHandlers/AI/Stability/StabilityTextToImagePromptHandler.cs b/BotNet.CommandHandlers/AI/Stability/StabilityTextToImagePromptHandler.cs
@@ -1,7 +1,6 @@
 using BotNet.Commands;
 using BotNet.Commands.AI.Stability;
 using BotNet.Commands.BotUpdate.Message;
-using BotNet.Services.RateLimit;
 using BotNet.Services.Stability.Models;
 using BotNet.Services.Stability.Skills;
 using Telegram.Bot;
@@ -14,25 +13,11 @@ public sealed class StabilityTextToImagePromptHandler(
 		ImageGenerationBot imageGenerationBot,
 		ITelegramMessageCache telegramMessageCache
 	) : ICommandHandler<StabilityTextToImagePrompt> {
-		internal static readonly RateLimiter IMAGE_GENERATION_RATE_LIMITER = RateLimiter.PerUser(1, TimeSpan.FromMinutes(5));
-
 		private readonly ITelegramBotClient _telegramBotClient = telegramBotClient;
 		private readonly ImageGenerationBot _imageGenerationBot = imageGenerationBot;
 		private readonly ITelegramMessageCache _telegramMessageCache = telegramMessageCache;
 
 		public Task Handle(StabilityTextToImagePrompt command, CancellationToken cancellationToken) {
-			try {
-				IMAGE_GENERATION_RATE_LIMITER.ValidateActionRate(command.ChatId, command.SenderId);
-			} catch (RateLimitExceededException exc) {
-				return _telegramBotClient.SendTextMessageAsync(
-					chatId: command.ChatId,
-					text: $"Anda belum mendapat giliran. Coba lagi {exc.Cooldown}.",
-					parseMode: ParseMode.Html,
-					replyToMessageId: command.PromptMessageId,
-					cancellationToken: cancellationToken
-				);
-			}
-
 			// Fire and forget
 			Task.Run(async () => {
 				try {

diff --git a/BotNet.CommandHandlers/Art/ArtCommandHandler.cs b/BotNet.CommandHandlers/Art/ArtCommandHandler.cs
@@ -0,0 +1,101 @@
+using BotNet.Commands;
+using BotNet.Commands.AI.OpenAI;
+using BotNet.Commands.AI.Stability;
+using BotNet.Commands.Art;
+using BotNet.Commands.CommandPrioritization;
+using BotNet.Services.MarkdownV2;
+using BotNet.Services.RateLimit;
+using Telegram.Bot;
+using Telegram.Bot.Types;
+using Telegram.Bot.Types.Enums;
+
+namespace BotNet.CommandHandlers.Art {
+	public sealed class ArtCommandHandler(
+		ITelegramBotClient telegramBotClient,
+		ICommandQueue commandQueue
+	) : ICommandHandler<ArtCommand> {
+		internal static readonly RateLimiter IMAGE_GENERATION_RATE_LIMITER = RateLimiter.PerUser(1, TimeSpan.FromMinutes(5));
+
+		private readonly ITelegramBotClient _telegramBotClient = telegramBotClient;
+		private readonly ICommandQueue _commandQueue = commandQueue;
+
+		public Task Handle(ArtCommand command, CancellationToken cancellationToken) {
+			try {
+				IMAGE_GENERATION_RATE_LIMITER.ValidateActionRate(command.ChatId, command.SenderId);
+			} catch (RateLimitExceededException exc) {
+				return _telegramBotClient.SendTextMessageAsync(
+					chatId: command.ChatId,
+					text: $"Anda belum mendapat giliran. Coba lagi {exc.Cooldown}.",
+					parseMode: ParseMode.Html,
+					replyToMessageId: command.PromptMessageId,
+					cancellationToken: cancellationToken
+				);
+			}
+
+			// Fire and forget
+			Task.Run(async () => {
+				try {
+					switch (command.CommandPriority) {
+						case CommandPriority.VIPChat: {
+								Message busyMessage = await _telegramBotClient.SendTextMessageAsync(
+									chatId: command.ChatId,
+									text: "Generating image… ⏳",
+									parseMode: ParseMode.Markdown,
+									replyToMessageId: command.PromptMessageId,
+									cancellationToken: cancellationToken
+								);
+
+								await _commandQueue.DispatchAsync(
+									new OpenAIImageGenerationPrompt(
+										callSign: "AI",
+										prompt: command.Prompt,
+										promptMessageId: command.PromptMessageId,
+										responseMessageId: busyMessage.MessageId,
+										chatId: command.ChatId,
+										senderId: command.SenderId,
+										commandPriority: command.CommandPriority
+									)
+								);
+							}
+							break;
+						case CommandPriority.HomeGroupChat: {
+								Message busyMessage = await _telegramBotClient.SendTextMessageAsync(
+									chatId: command.ChatId,
+									text: "Generating image… ⏳",
+									parseMode: ParseMode.Markdown,
+									replyToMessageId: command.PromptMessageId,
+									cancellationToken: cancellationToken
+								);
+
+								await _commandQueue.DispatchAsync(
+									new StabilityTextToImagePrompt(
+										callSign: "AI",
+										prompt: command.Prompt,
+										promptMessageId: command.PromptMessageId,
+										responseMessageId: busyMessage.MessageId,
+										chatId: command.ChatId,
+										senderId: command.SenderId,
+										commandPriority: command.CommandPriority
+									)
+								);
+							}
+							break;
+						default:
+							await _telegramBotClient.SendTextMessageAsync(
+								chatId: command.ChatId,
+								text: MarkdownV2Sanitizer.Sanitize("Image generation tidak bisa dipakai di sini."),
+								parseMode: ParseMode.MarkdownV2,
+								replyToMessageId: command.PromptMessageId,
+								cancellationToken: cancellationToken
+							);
+							break;
+					}
+				} catch (OperationCanceledException) {
+					// Terminate gracefully
+				}
+			});
+
+			return Task.CompletedTask;
+		}
+	}
+}
diff --git a/BotNet.CommandHandlers/BotUpdate/Message/SlashCommandHandler.cs b/BotNet.CommandHandlers/BotUpdate/Message/SlashCommandHandler.cs
@@ -1,6 +1,6 @@
 using BotNet.Commands;
 using BotNet.Commands.AI.OpenAI;
-using BotNet.Commands.AI.Stability;
+using BotNet.Commands.Art;
 using BotNet.Commands.BMKG;
 using BotNet.Commands.BotUpdate.Message;
 using BotNet.Commands.Common;

diff --git a/BotNet.Commands/AI/Stability/ArtCommand.cs → BotNet.Commands/Art/ArtCommand.cs b/BotNet.Commands/AI/Stability/ArtCommand.cs → BotNet.Commands/Art/ArtCommand.cs
@@ -3,7 +3,7 @@
 using BotNet.Commands.Common;
 using Telegram.Bot.Types.Enums;
 
-namespace BotNet.Commands.AI.Stability {
+namespace BotNet.Commands.Art {
 	public sealed record ArtCommand : ICommand {
 		public string Prompt { get; }
 		public int PromptMessageId { get; }