From 3ac6aecc76dc5dc2d90af74a22e14ac67f8aab25 Mon Sep 17 00:00:00 2001
From: Ronny Gunawan <3048897+ronnygunawan@users.noreply.github.com>
Date: Mon, 4 Dec 2023 22:15:23 +0700
Subject: [PATCH] Handle Stability errors

---
 BotNet.Services/BotCommands/Art.cs            |  8 ++++
 BotNet.Services/BotCommands/OpenAI.cs         | 47 +++++++++++++------
 .../Models/ContentFilteredException.cs        |  5 ++
 .../Stability/Models/TextToImageResponse.cs   |  3 +-
 BotNet.Services/Stability/StabilityClient.cs  |  4 ++
 5 files changed, 50 insertions(+), 17 deletions(-)
 create mode 100644 BotNet.Services/Stability/Models/ContentFilteredException.cs

diff --git a/BotNet.Services/BotCommands/Art.cs b/BotNet.Services/BotCommands/Art.cs
index a4dc82d..77f7bea 100644
--- a/BotNet.Services/BotCommands/Art.cs
+++ b/BotNet.Services/BotCommands/Art.cs
@@ -4,6 +4,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using BotNet.Services.RateLimit;
+using BotNet.Services.Stability.Models;
 using BotNet.Services.ThisXDoesNotExist;
 using Microsoft.Extensions.DependencyInjection;
 using Telegram.Bot;
@@ -50,6 +51,13 @@ await botClient.SendPhotoAsync(
 								photo: new InputFileStream(imageStream, "art.png"),
 								replyToMessageId: message.MessageId,
 								cancellationToken: cancellationToken);
+						} catch (ContentFilteredException) {
+							await botClient.SendTextMessageAsync(
+								chatId: message.Chat.Id,
+								text: "<code>Content filtered</code>",
+								parseMode: ParseMode.Html,
+								replyToMessageId: message.MessageId,
+								cancellationToken: cancellationToken);
 						} catch {
 							await botClient.SendTextMessageAsync(
 								chatId: message.Chat.Id,
diff --git a/BotNet.Services/BotCommands/OpenAI.cs b/BotNet.Services/BotCommands/OpenAI.cs
index 510894c..72569fd 100644
--- a/BotNet.Services/BotCommands/OpenAI.cs
+++ b/BotNet.Services/BotCommands/OpenAI.cs
@@ -12,6 +12,7 @@
 using BotNet.Services.OpenAI.Models;
 using BotNet.Services.OpenAI.Skills;
 using BotNet.Services.RateLimit;
+using BotNet.Services.Stability.Models;
 using BotNet.Services.Stability.Skills;
 using Microsoft.Extensions.DependencyInjection;
 using RG.Ninja;
@@ -851,26 +852,35 @@ await serviceProvider.GetRequiredService<FriendlyBot>().StreamChatAsync(
 								//	prompt: message.Text!,
 								//	cancellationToken: cancellationToken
 								//);
-								byte[] generatedImage = await serviceProvider.GetRequiredService<Stability.Skills.ImageGenerationBot>().GenerateImageAsync(
-									prompt: message.Text!,
-									cancellationToken: cancellationToken
-								);
-								using MemoryStream generatedImageStream = new(generatedImage);
 								try {
-									await botClient.DeleteMessageAsync(
+									byte[] generatedImage = await serviceProvider.GetRequiredService<Stability.Skills.ImageGenerationBot>().GenerateImageAsync(
+										prompt: message.Text!,
+										cancellationToken: cancellationToken
+									);
+									using MemoryStream generatedImageStream = new(generatedImage);
+									try {
+										await botClient.DeleteMessageAsync(
+											chatId: busyMessage.Chat.Id,
+											messageId: busyMessage.MessageId,
+											cancellationToken: cancellationToken
+										);
+									} catch (OperationCanceledException) {
+										throw;
+									}
+									await botClient.SendPhotoAsync(
+										chatId: message.Chat.Id,
+										photo: new InputFileStream(generatedImageStream, "art.png"),
+										replyToMessageId: message.MessageId,
+										cancellationToken: cancellationToken
+									);
+								} catch (ContentFilteredException) {
+									await botClient.EditMessageTextAsync(
 										chatId: busyMessage.Chat.Id,
 										messageId: busyMessage.MessageId,
-										cancellationToken: cancellationToken
+										text: "<code>Content filtered.</code>",
+										parseMode: ParseMode.Html
 									);
-								} catch (OperationCanceledException) {
-									throw;
 								}
-								await botClient.SendPhotoAsync(
-									chatId: message.Chat.Id,
-									photo: new InputFileStream(generatedImageStream, "art.png"),
-									replyToMessageId: message.MessageId,
-									cancellationToken: cancellationToken
-								);
 								break;
 							}
 					}
@@ -901,6 +911,13 @@ await botClient.SendTextMessageAsync(
 					parseMode: ParseMode.Html,
 					replyToMessageId: message.MessageId,
 					cancellationToken: cancellationToken);
+			} catch (HttpRequestException exc) {
+				await botClient.SendTextMessageAsync(
+					chatId: message.Chat.Id,
+					text: "<code>Unknown error.</code>",
+					parseMode: ParseMode.Html,
+					replyToMessageId: message.MessageId,
+					cancellationToken: cancellationToken);
 			} catch (OperationCanceledException) {
 				await botClient.SendTextMessageAsync(
 					chatId: message.Chat.Id,
diff --git a/BotNet.Services/Stability/Models/ContentFilteredException.cs b/BotNet.Services/Stability/Models/ContentFilteredException.cs
new file mode 100644
index 0000000..29cd174
--- /dev/null
+++ b/BotNet.Services/Stability/Models/ContentFilteredException.cs
@@ -0,0 +1,5 @@
+using System;
+
+namespace BotNet.Services.Stability.Models {
+	public sealed class ContentFilteredException : Exception { }
+}
diff --git a/BotNet.Services/Stability/Models/TextToImageResponse.cs b/BotNet.Services/Stability/Models/TextToImageResponse.cs
index dd36981..6eb79d7 100644
--- a/BotNet.Services/Stability/Models/TextToImageResponse.cs
+++ b/BotNet.Services/Stability/Models/TextToImageResponse.cs
@@ -7,7 +7,6 @@ List<Artifact> Artifacts
 
 	internal sealed record Artifact(
 		string Base64,
-		string FinishReason,
-		int Seed
+		string FinishReason
 	);
 }
diff --git a/BotNet.Services/Stability/StabilityClient.cs b/BotNet.Services/Stability/StabilityClient.cs
index e1c6f57..f8730cc 100644
--- a/BotNet.Services/Stability/StabilityClient.cs
+++ b/BotNet.Services/Stability/StabilityClient.cs
@@ -67,6 +67,10 @@ CancellationToken cancellationToken
 
 			TextToImageResponse? responseData = JsonSerializer.Deserialize<TextToImageResponse>(responseJson, CAMEL_CASE_SERIALIZER_OPTIONS);
 
+			if (responseData is { Artifacts: [Artifact { FinishReason: "CONTENT_FILTERED" }] }) {
+				throw new ContentFilteredException();
+			}
+
 			if (responseData is not { Artifacts: [Artifact { FinishReason: "SUCCESS", Base64: var base64 }] }) {
 				throw new HttpRequestException();
 			}