Refactor GptImage and TransformImage command handlers

Author: selfdocumentingcode

src/Kattbot/CommandHandlers/Images/GptImage.cs

@@ -28,7 +28,37 @@ public GptImageCommand(CommandContext ctx, string prompt)
     public string Prompt { get; }
 }
 
-public class GptImageHandler : IRequestHandler<GptImageCommand>
+public class GptImageAvatarCommand : CommandRequest
+{
+    public GptImageAvatarCommand(CommandContext ctx, DiscordUser user, string prompt)
+        : base(ctx)
+    {
+        Prompt = prompt;
+        User = user;
+    }
+
+    public string Prompt { get; }
+
+    public DiscordUser User { get; }
+}
+
+public class GptImageEmoteCommand : CommandRequest
+{
+    public GptImageEmoteCommand(CommandContext ctx, DiscordEmoji emoji, string prompt)
+        : base(ctx)
+    {
+        Prompt = prompt;
+        Emoji = emoji;
+    }
+
+    public string Prompt { get; }
+
+    public DiscordEmoji Emoji { get; }
+}
+
+public class GptImageHandler : IRequestHandler<GptImageCommand>,
+    IRequestHandler<GptImageAvatarCommand>,
+    IRequestHandler<GptImageEmoteCommand>
 {
     private const string GptImageModel = "gpt-image-1";
     private const string Moderation = "low";
@@ -38,17 +68,20 @@ public class GptImageHandler : IRequestHandler<GptImageCommand>
 
     private readonly GptImagesHttpClient _gptImagesHttpClient;
     private readonly ImageService _imageService;
+    private readonly DiscordResolver _discordResolver;
     private readonly ILogger<GptImageHandler> _logger;
 
     private readonly string[] _supportedImageFormats = ["png", "jpg", "webp"];
 
     public GptImageHandler(
         GptImagesHttpClient gptImagesHttpClient,
         ImageService imageService,
+        DiscordResolver discordResolver,
         ILogger<GptImageHandler> logger)
     {
         _gptImagesHttpClient = gptImagesHttpClient;
         _imageService = imageService;
+        _discordResolver = discordResolver;
         _logger = logger;
     }
 
@@ -66,6 +99,7 @@ public async Task Handle(GptImageCommand request, CancellationToken cancellation
             string? imageUrl = await discordMessage.GetImageUrlFromMessage(_logger);
 
             CreateImageResponse response;
+            var userId = request.Ctx.User.Id.ToString();
 
             if (imageUrl == null)
             {
@@ -75,40 +109,68 @@ public async Task Handle(GptImageCommand request, CancellationToken cancellation
                     Model = GptImageModel,
                     Quality = Quality,
                     Moderation = Moderation,
-                    User = request.Ctx.User.Id.ToString(),
+                    User = userId,
                 };
 
                 response = await _gptImagesHttpClient.CreateImage(imageRequest, cancellationToken);
             }
             else
             {
-                Image editImage = await _imageService.DownloadImage(imageUrl);
+                response = await CreateImageEdit(imageUrl, prompt, userId, cancellationToken);
+            }
 
-                Image imageInSupportedFormat =
-                    await ImageService.EnsureSupportedImageFormatOrPng(editImage, _supportedImageFormats);
+            if (response.Data == null || !response.Data.Any()) throw new Exception("Empty result");
 
-                Image resizedImage = await ImageService.EnsureMaxImageFileSize(
-                    imageInSupportedFormat,
-                    MaxImageSizeInMb);
+            ImageResponseData imageData = response.Data.First();
 
-                var tempFileName = $"{Guid.NewGuid()}.png";
+            Image image = ImageService.ConvertBase64ToImage(imageData.B64Json);
 
-                ImageStreamResult inputImageStream = await ImageService.GetImageStream(resizedImage);
+            ImageStreamResult imageStream = await ImageService.GetImageStream(image);
 
-                var editImageRequest = new CreateImageEditRequest
-                {
-                    Prompt = prompt,
-                    Image = inputImageStream.MemoryStream.ToArray(),
-                    Model = GptImageModel,
-                    Quality = Quality,
-                    User = request.Ctx.User.Id.ToString(),
-                };
+            string fileName = prompt.ToSafeFilename(imageStream.FileExtension);
 
-                response = await _gptImagesHttpClient.CreateImageEdit(
-                    editImageRequest,
-                    tempFileName,
-                    cancellationToken);
-            }
+            string truncatedPrompt = prompt.Length > DiscordConstants.MaxEmbedTitleLength
+                ? $"{prompt[..(DiscordConstants.MaxEmbedTitleLength - 3)]}..."
+                : prompt;
+
+            DiscordEmbedBuilder eb = new DiscordEmbedBuilder()
+                .WithTitle(truncatedPrompt)
+                .WithImageUrl($"attachment://{fileName}");
+
+            DiscordMessageBuilder mb = new DiscordMessageBuilder()
+                .AddFile(fileName, imageStream.MemoryStream)
+                .AddEmbed(eb)
+                .WithContent($"There you go {request.Ctx.Member?.Mention ?? "Unknown user"}");
+
+            await request.Ctx.RespondAsync(mb);
+        }
+        finally
+        {
+            await ackMessage.DeleteAsync();
+        }
+    }
+
+    public async Task Handle(GptImageAvatarCommand request, CancellationToken cancellationToken)
+    {
+        CommandContext ctx = request.Ctx;
+        DiscordMessage discordMessage = ctx.Message;
+        DiscordUser user = request.User;
+        DiscordGuild guild = ctx.Guild;
+
+        DiscordMessage ackMessage = await request.Ctx.RespondAsync("Working on it");
+
+        try
+        {
+            string prompt = discordMessage.SubstituteMentions(request.Prompt);
+
+            DiscordMember userAsMember = await _discordResolver.ResolveGuildMember(guild, user.Id) ??
+                                         throw new Exception("Invalid user");
+
+            string imageUrl = userAsMember.GuildAvatarUrl ?? userAsMember.AvatarUrl;
+
+            var userId = request.Ctx.User.Id.ToString();
+
+            CreateImageResponse response = await CreateImageEdit(imageUrl, prompt, userId, cancellationToken);
 
             if (response.Data == null || !response.Data.Any()) throw new Exception("Empty result");
 
@@ -140,4 +202,90 @@ public async Task Handle(GptImageCommand request, CancellationToken cancellation
             await ackMessage.DeleteAsync();
         }
     }
+
+    public async Task Handle(GptImageEmoteCommand request, CancellationToken cancellationToken)
+    {
+        CommandContext ctx = request.Ctx;
+        DiscordMessage discordMessage = ctx.Message;
+
+        DiscordEmoji emoji = request.Emoji;
+
+        DiscordMessage ackMessage = await request.Ctx.RespondAsync("Working on it");
+
+        try
+        {
+            string prompt = discordMessage.SubstituteMentions(request.Prompt);
+
+            string imageUrl = emoji.GetEmojiImageUrl();
+
+            var userId = request.Ctx.User.Id.ToString();
+
+            CreateImageResponse response = await CreateImageEdit(imageUrl, prompt, userId, cancellationToken);
+
+            if (response.Data == null || !response.Data.Any()) throw new Exception("Empty result");
+
+            ImageResponseData imageData = response.Data.First();
+
+            Image image = ImageService.ConvertBase64ToImage(imageData.B64Json);
+
+            ImageStreamResult imageStream = await ImageService.GetImageStream(image);
+
+            string fileName = prompt.ToSafeFilename(imageStream.FileExtension);
+
+            string truncatedPrompt = prompt.Length > DiscordConstants.MaxEmbedTitleLength
+                ? $"{prompt[..(DiscordConstants.MaxEmbedTitleLength - 3)]}..."
+                : prompt;
+
+            DiscordEmbedBuilder eb = new DiscordEmbedBuilder()
+                .WithTitle(truncatedPrompt)
+                .WithImageUrl($"attachment://{fileName}");
+
+            DiscordMessageBuilder mb = new DiscordMessageBuilder()
+                .AddFile(fileName, imageStream.MemoryStream)
+                .AddEmbed(eb)
+                .WithContent($"There you go {request.Ctx.Member?.Mention ?? "Unknown user"}");
+
+            await request.Ctx.RespondAsync(mb);
+        }
+        finally
+        {
+            await ackMessage.DeleteAsync();
+        }
+    }
+
+    private async Task<CreateImageResponse> CreateImageEdit(
+        string imageUrl,
+        string prompt,
+        string userId,
+        CancellationToken cancellationToken)
+    {
+        Image editImage = await _imageService.DownloadImage(imageUrl);
+
+        Image imageInSupportedFormat =
+            await ImageService.EnsureSupportedImageFormatOrPng(editImage, _supportedImageFormats);
+
+        Image resizedImage = await ImageService.EnsureMaxImageFileSize(
+            imageInSupportedFormat,
+            MaxImageSizeInMb);
+
+        var tempFileName = $"{Guid.NewGuid()}.png";
+
+        ImageStreamResult inputImageStream = await ImageService.GetImageStream(resizedImage);
+
+        var editImageRequest = new CreateImageEditRequest
+        {
+            Prompt = prompt,
+            Image = inputImageStream.MemoryStream.ToArray(),
+            Model = GptImageModel,
+            Quality = Quality,
+            User = userId,
+        };
+
+        CreateImageResponse response = await _gptImagesHttpClient.CreateImageEdit(
+            editImageRequest,
+            tempFileName,
+            cancellationToken);
+
+        return response;
+    }
 }