Vision bot can now look at stickers

teknologi-umum · Dec 3, 2023 · 18fd9d0 · 18fd9d0
1 parent 81a81c1
commit 18fd9d0
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 18 deletions.
diff --git a/BotNet.Services/BotCommands/OpenAI.cs b/BotNet.Services/BotCommands/OpenAI.cs
@@ -775,21 +775,26 @@ CancellationToken cancellationToken
 					: CHAT_GROUP_RATE_LIMITER
 				).ValidateActionRate(message.Chat.Id, message.From!.Id);
 
-				PhotoSize? photoSize;
-				string? caption;
+				string? fileId;
+				string? prompt;
 				if (message is { Photo.Length: > 0, Caption: { } }) {
-					photoSize = message.Photo.OrderByDescending(photoSize => photoSize.Width).First();
-					caption = message.Caption;
-				} else if (message.ReplyToMessage is { Photo.Length: > 0, Caption: { } }) {
-					photoSize = message.ReplyToMessage.Photo.OrderByDescending(photoSize => photoSize.Width).First();
-					caption = message.ReplyToMessage.Caption;
+					fileId = message.Photo.OrderByDescending(photoSize => photoSize.Width).First().FileId;
+					prompt = message.Caption;
+				} else if (message.ReplyToMessage is { Photo.Length: > 0 }
+					&& message.Text is { }) {
+					fileId = message.ReplyToMessage.Photo.OrderByDescending(photoSize => photoSize.Width).First().FileId;
+					prompt = message.Caption;
+				} else if (message.ReplyToMessage is { Sticker: { } }
+					&& message.Text is { }) {
+					fileId = message.ReplyToMessage.Sticker.FileId;
+					prompt = message.Caption;
 				} else {
-					photoSize = null;
-					caption = null;
+					fileId = null;
+					prompt = null;
 				}
 
-				if (photoSize != null && caption != null) {
-					(string? imageBase64, string? error) = await GetImageBase64Async(botClient, photoSize, cancellationToken);
+				if (fileId != null && prompt != null) {
+					(string? imageBase64, string? error) = await GetImageBase64Async(botClient, fileId, cancellationToken);
 					if (error != null) {
 						await botClient.SendTextMessageAsync(
 							chatId: message.Chat.Id,
@@ -800,7 +805,7 @@ await botClient.SendTextMessageAsync(
 						return;
 					}
 					await serviceProvider.GetRequiredService<VisionBot>().StreamChatAsync(
-						message: caption,
+						message: prompt,
 						imageBase64: imageBase64!,
 						chatId: message.Chat.Id,
 						replyToMessageId: message.MessageId
@@ -882,11 +887,11 @@ await botClient.SendTextMessageAsync(
 			}
 		}
 
-		private static async Task<(string? ImageBase64, string? Error)> GetImageBase64Async(ITelegramBotClient botClient, PhotoSize photoSize, CancellationToken cancellationToken) {
+		private static async Task<(string? ImageBase64, string? Error)> GetImageBase64Async(ITelegramBotClient botClient, string fileId, CancellationToken cancellationToken) {
 			// Download photo
 			using MemoryStream originalImageStream = new();
 			await botClient.GetInfoAndDownloadFileAsync(
-				fileId: photoSize.FileId,
+				fileId: fileId,
 				destination: originalImageStream,
 				cancellationToken: cancellationToken);
 			byte[] originalImage = originalImageStream.ToArray();

diff --git a/BotNet.Services/OpenAI/OpenAIClient.cs b/BotNet.Services/OpenAI/OpenAIClient.cs
@@ -127,13 +127,36 @@ [EnumeratorCancellation] CancellationToken cancellationToken
 
 			while (!streamReader.EndOfStream) {
 				string? line = await streamReader.ReadLineAsync(cancellationToken);
-				if (line == null) break;
+
+				if (line == null) {
+					yield return (
+						Result: result.ToString(),
+						Stop: true
+					);
+					yield break;
+				}
+
 				if (line == "") continue;
 				if (!line.StartsWith("data: ", out string? json)) break;
-				if (json == "[DONE]") break;
+
+				if (json == "[DONE]") {
+					yield return (
+						Result: result.ToString(),
+						Stop: true
+					);
+					yield break;
+				}
+
 				CompletionResult? completionResult = JsonSerializer.Deserialize<CompletionResult>(json, JSON_SERIALIZER_OPTIONS);
-				if (completionResult == null) break;
-				if (completionResult.Choices.Count == 0) break;
+
+				if (completionResult == null || completionResult.Choices.Count == 0) {
+					yield return (
+						Result: result.ToString(),
+						Stop: true
+					);
+					yield break;
+				}
+
 				result.Append(completionResult.Choices[0].Delta!.Content);
 
 				if (completionResult.Choices[0].FinishReason == "stop") {