Merge pull request #522 from iceljc/features/refine-image-generation

Features/refine image generation

src/Infrastructure/BotSharp.Abstraction/Conversations/Models/RoleDialogModel.cs

@@ -1,6 +1,7 @@
 using BotSharp.Abstraction.Functions.Models;
 using BotSharp.Abstraction.Messaging;
 using BotSharp.Abstraction.Messaging.Models.RichContent;
+using BotSharp.Abstraction.MLTasks;
 
 namespace BotSharp.Abstraction.Conversations.Models;
 
@@ -87,6 +88,13 @@ public class RoleDialogModel : ITrackableMessage
 
     public List<BotSharpFile> Files { get; set; } = new List<BotSharpFile>();
 
+    /// <summary>
+    /// The images generated by AI
+    /// </summary>
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    [JsonPropertyName("generated_images")]
+    public List<ImageGeneration> GeneratedImages { get; set; } = new List<ImageGeneration>();
+
     private RoleDialogModel()
     {
     }

src/Infrastructure/BotSharp.Abstraction/Files/IBotSharpFileService.cs

@@ -23,6 +23,14 @@ public interface IBotSharpFileService
     bool DeleteMessageFiles(string conversationId, IEnumerable<string> messageIds, string targetMessageId, string? newMessageId = null);
     bool DeleteConversationFiles(IEnumerable<string> conversationIds);
 
+    /// <summary>
+    /// Take screenshots of pdf pages and get response from llm
+    /// </summary>
+    /// <param name="prompt"></param>
+    /// <param name="files">Pdf files</param>
+    /// <returns></returns>
+    Task<string> InstructPdf(string? provider, string? model, string? modelId, string prompt, List<BotSharpFile> files);
+
     /// <summary>
     /// Get file bytes and content type from data, e.g., "data:image/png;base64,aaaaaaaaa"
     /// </summary>

src/Infrastructure/BotSharp.Abstraction/Files/Models/ImageGeneration.cs

@@ -0,0 +1,13 @@
+namespace BotSharp.Abstraction.Files.Models;
+
+public class ImageGeneration
+{
+    [JsonPropertyName("image_url")]
+    public string? ImageUrl { get; set; }
+
+    [JsonPropertyName("image_data")]
+    public string? ImageData { get; set; }
+
+    [JsonPropertyName("description")]
+    public string Description { get; set; } = string.Empty;
+}

src/Infrastructure/BotSharp.Core/BotSharp.Core.csproj

@@ -180,10 +180,8 @@
     <PackageReference Include="EntityFrameworkCore.BootKit" Version="8.4.2" />
     <PackageReference Include="Fluid.Core" Version="2.8.0" />
     <PackageReference Include="Microsoft.AspNetCore.StaticFiles" Version="2.2.0" />
+    <PackageReference Include="Microsoft.Extensions.Http" Version="8.0.0" />
     <PackageReference Include="Nanoid" Version="3.0.0" />
-    <PackageReference Include="PdfiumViewer" Version="2.13.0" />
-    <PackageReference Include="PdfiumViewer.Native.x86.v8-xfa" Version="2018.4.8.256" />
-    <PackageReference Include="PdfiumViewer.Native.x86_64.v8-xfa" Version="2018.4.8.256" />
     <PackageReference Include="RedLock.net" Version="2.3.2" />
     <PackageReference Include="System.Drawing.Common" Version="8.0.6" />
   </ItemGroup>

src/Infrastructure/BotSharp.Core/Files/Converters/PdfiumConverter.cs

@@ -1,39 +0,0 @@
-using BotSharp.Abstraction.Files.Converters;
-using PdfiumViewer;
-using System.IO;
-
-namespace BotSharp.Core.Files.Converters;
-
-public class PdfiumConverter : IPdf2ImageConverter
-{
-    public async Task<IEnumerable<string>> ConvertPdfToImages(string pdfLocation, string imageFolderLocation)
-    {
-        var paths = new List<string>();
-        if (string.IsNullOrWhiteSpace(imageFolderLocation)) return paths;
-
-        if (Directory.Exists(imageFolderLocation))
-        {
-            Directory.Delete(imageFolderLocation, true);
-        }
-        Directory.CreateDirectory(imageFolderLocation);
-
-        var guid = Guid.NewGuid().ToString();
-        using (var document = PdfDocument.Load(pdfLocation))
-        {
-            var pages = document.PageCount;
-
-            for (var page = 0; page < pages; page++)
-            {
-                var size = document.PageSizes[page];
-                using (var image = document.Render(page, (int)size.Width, (int)size.Height, 96, 96, true))
-                {
-                    var imagePath = Path.Combine(imageFolderLocation, $"{guid}_pg_{page + 1}.png");
-                    image.Save(imagePath, System.Drawing.Imaging.ImageFormat.Png);
-                    paths.Add(imagePath);
-                }
-            }
-        }
-
-        return await Task.FromResult(paths);
-    }
-}

src/Infrastructure/BotSharp.Core/Files/FilePlugin.cs

@@ -1,6 +1,5 @@
-using BotSharp.Abstraction.Files.Converters;
-using BotSharp.Core.Files.Converters;
 using BotSharp.Core.Files.Hooks;
+using BotSharp.Core.Files.Services;
 using Microsoft.Extensions.Configuration;
 
 namespace BotSharp.Core.Files;
@@ -20,6 +19,5 @@ public class FilePlugin : IBotSharpPlugin
 
         services.AddScoped<IAgentHook, FileAnalyzerHook>();
         services.AddScoped<IAgentToolHook, FileAnalyzerToolHook>();
-        services.AddScoped<IPdf2ImageConverter, PdfiumConverter>();
     }
 }

src/Infrastructure/BotSharp.Core/Files/Services/BotSharpFileService.Conversation.cs

@@ -1,11 +1,9 @@
 using BotSharp.Abstraction.Files.Converters;
-using BotSharp.Core.Files.Converters;
 using Microsoft.EntityFrameworkCore;
 using System.IO;
-using System.Linq;
 using System.Threading;
 
-namespace BotSharp.Core.Files;
+namespace BotSharp.Core.Files.Services;
 
 public partial class BotSharpFileService
 {
@@ -118,7 +116,7 @@ public partial class BotSharpFileService
         }
         catch (Exception ex)
         {
-            _logger.LogWarning($"Error when reading conversation ({conversationId}) files: {ex.Message}");
+            _logger.LogWarning($"Error when reading conversation ({conversationId}) files: {ex.Message}\r\n{ex.InnerException}\r\n{ex.StackTrace}");
         }
 
         return files;
@@ -281,7 +279,7 @@ public partial class BotSharpFileService
 
         foreach (var conversationId in conversationIds)
         {
-            var convDir = FindConversationDirectory(conversationId);
+            var convDir = GetConversationDirectory(conversationId);
             if (!ExistDirectory(convDir)) continue;
 
             Directory.Delete(convDir, true);
@@ -305,7 +303,7 @@ public partial class BotSharpFileService
         return dir;
     }
 
-    private string? FindConversationDirectory(string conversationId)
+    private string? GetConversationDirectory(string conversationId)
     {
         if (string.IsNullOrEmpty(conversationId)) return null;
 
@@ -318,14 +316,18 @@ public partial class BotSharpFileService
         var converters = _services.GetServices<IPdf2ImageConverter>();
         if (converters.IsNullOrEmpty()) return Enumerable.Empty<string>();
 
-        var converter = converters.FirstOrDefault(x => x.GetType().Name != typeof(PdfiumConverter).Name);
+        var converter = GetPdf2ImageConverter();
         if (converter == null)
         {
-            converter = converters.FirstOrDefault(x => x.GetType().Name == typeof(PdfiumConverter).Name);
-            if (converter == null) return Enumerable.Empty<string>();
+            return Enumerable.Empty<string>();
         }
-
         return await converter.ConvertPdfToImages(pdfLoc, imageLoc);
     }
+
+    private IPdf2ImageConverter? GetPdf2ImageConverter()
+    {
+        var converters = _services.GetServices<IPdf2ImageConverter>();
+        return converters.FirstOrDefault();
+    }
     #endregion
 }

src/Infrastructure/BotSharp.Core/Files/Services/BotSharpFileService.Pdf.cs

@@ -0,0 +1,146 @@
+using System.IO;
+using System.Linq;
+using System.Net.Http;
+using System.Threading;
+
+namespace BotSharp.Core.Files.Services;
+
+public partial class BotSharpFileService
+{
+    public async Task<string> InstructPdf(string? provider, string? model, string? modelId, string prompt, List<BotSharpFile> files)
+    {
+        var content = string.Empty;
+
+        if (string.IsNullOrWhiteSpace(prompt) || files.IsNullOrEmpty())
+        {
+            return content;
+        }
+
+        var guid = Guid.NewGuid().ToString();
+        var sessionDir = GetSessionDirectory(guid);
+        if (!ExistDirectory(sessionDir))
+        {
+            Directory.CreateDirectory(sessionDir);
+        }
+
+        try
+        {
+            var pdfFiles = await SaveFiles(sessionDir, files);
+            var images = await ConvertPdfToImages(pdfFiles);
+            if (images.IsNullOrEmpty()) return content;
+
+            var completion = CompletionProvider.GetChatCompletion(_services, provider: provider ?? "openai",
+                model: model, modelId: modelId ?? "gpt-4", multiModal: true);
+            var message = await completion.GetChatCompletions(new Agent()
+            {
+                Id = Guid.Empty.ToString(),
+            }, new List<RoleDialogModel>
+            {
+                new RoleDialogModel(AgentRole.User, prompt)
+                {
+                    Files = images.Select(x => new BotSharpFile { FileStorageUrl = x }).ToList()
+                }
+            });
+
+            content = message.Content;
+            return content;
+        }
+        catch (Exception ex)
+        {
+            _logger.LogError($"Error when analyzing pdf in file service: {ex.Message}");
+            return content;
+        }
+        finally
+        {
+            Directory.Delete(sessionDir, true);
+        }
+    }
+
+    #region Private methods
+    private string GetSessionDirectory(string id)
+    {
+        var dir = Path.Combine(_baseDir, SESSION_FOLDER, id);
+        return dir;
+    }
+
+    private async Task<IEnumerable<string>> SaveFiles(string dir, List<BotSharpFile> files, string extension = "pdf")
+    {
+        if (string.IsNullOrWhiteSpace(dir) || files.IsNullOrEmpty())
+        {
+            return Enumerable.Empty<string>();
+        }
+
+        var locs = new List<string>();
+        foreach (var file in files)
+        {
+            try
+            {
+                var bytes = new byte[0];
+                if (!string.IsNullOrEmpty(file.FileUrl))
+                {
+                    var http = _services.GetRequiredService<IHttpClientFactory>();
+                    using var client = http.CreateClient();
+                    bytes = await client.GetByteArrayAsync(file.FileUrl);
+                }
+                else if (!string.IsNullOrEmpty(file.FileData))
+                {
+                    (_, bytes) = GetFileInfoFromData(file.FileData);
+                }
+
+                if (!bytes.IsNullOrEmpty())
+                {
+                    var guid = Guid.NewGuid().ToString();
+                    var fileDir = Path.Combine(dir, guid);
+                    if (!ExistDirectory(fileDir))
+                    {
+                        Directory.CreateDirectory(fileDir);
+                    }
+
+                    var pdfDir = Path.Combine(fileDir, $"{guid}.{extension}");
+                    using (var fs = new FileStream(pdfDir, FileMode.Create))
+                    {
+                        fs.Write(bytes, 0, bytes.Length);
+                        fs.Close();
+                        locs.Add(pdfDir);
+                        Thread.Sleep(100);
+                    }
+                }
+            }
+            catch (Exception ex)
+            {
+                _logger.LogWarning($"Error when saving pdf file: {ex.Message}");
+                continue;
+            }
+        }
+        return locs;
+    }
+
+    private async Task<IEnumerable<string>> ConvertPdfToImages(IEnumerable<string> files)
+    {
+        var images = new List<string>();
+        var converter = GetPdf2ImageConverter();
+        if (converter == null || files.IsNullOrEmpty())
+        {
+            return images;
+        }
+
+        foreach (var file in files)
+        {
+            try
+            {
+                var segs = file.Split(Path.DirectorySeparatorChar);
+                var dir = string.Join(Path.DirectorySeparatorChar, segs.SkipLast(1));
+                var folder = Path.Combine(dir, "screenshots");
+                var urls = await converter.ConvertPdfToImages(file, folder);
+                images.AddRange(urls);
+            }
+            catch (Exception ex)
+            {
+                _logger.LogWarning($"Error when converting pdf file to images ({file}): {ex.Message}");
+                continue;
+            }
+        }
+        return images;
+    }
+    #endregion
+}

src/Infrastructure/BotSharp.Core/Files/Services/BotSharpFileService.User.cs

@@ -1,6 +1,6 @@
 using System.IO;
 
-namespace BotSharp.Core.Files;
+namespace BotSharp.Core.Files.Services;
 
 public partial class BotSharpFileService
 {

src/Infrastructure/BotSharp.Core/Files/Services/BotSharpFileService.cs

@@ -1,9 +1,7 @@
 using Microsoft.AspNetCore.StaticFiles;
-using System;
 using System.IO;
-using System.Threading;
 
-namespace BotSharp.Core.Files;
+namespace BotSharp.Core.Files.Services;
 
 public partial class BotSharpFileService : IBotSharpFileService
 {
@@ -22,6 +20,7 @@ public partial class BotSharpFileService : IBotSharpFileService
     private const string BOT_FILE_FOLDER = "bot";
     private const string USERS_FOLDER = "users";
     private const string USER_AVATAR_FOLDER = "avatar";
+    private const string SESSION_FOLDER = "sessions";
 
     private const int MIN_OFFSET = 1;
     private const int MAX_OFFSET = 5;

src/Infrastructure/BotSharp.Core/Users/Services/UserService.cs

@@ -96,10 +96,13 @@ public class UserService : IUserService
         }
 
         //verify password is correct or not.
-        var hashPassword = Utilities.HashTextMd5($"{password}{record.Salt}");
-        if (hashPassword != record.Password)
+        if (record != null)
         {
-            return default;
+            var hashPassword = Utilities.HashTextMd5($"{password}{record.Salt}");
+            if (hashPassword != record.Password)
+            {
+                return default;
+            }
         }
 
         User? user = record;

src/Infrastructure/BotSharp.OpenAPI/Controllers/InstructModeController.cs

@@ -125,8 +125,8 @@ public class InstructModeController : ControllerBase
                 new RoleDialogModel(AgentRole.User, input.Text)
             });
             
-            imageViewModel.RevisedPrompt = message.Content;
-            imageViewModel.Data = message.Data;
+            imageViewModel.Content = message.Content;
+            imageViewModel.Images = message.GeneratedImages.Select(x => ImageViewModel.ToViewModel(x)).ToList();
             return imageViewModel;
         }
         catch (Exception ex)
@@ -137,4 +137,27 @@ public class InstructModeController : ControllerBase
             return imageViewModel;
         }
     }
+
+    [HttpPost("/instruct/pdf-completion")]
+    public async Task<PdfCompletionViewModel> PdfCompletion([FromBody] IncomingMessageModel input)
+    {
+        var state = _services.GetRequiredService<IConversationStateService>();
+        input.States.ForEach(x => state.SetState(x.Key, x.Value, activeRounds: x.ActiveRounds, source: StateSource.External));
+        var viewModel = new PdfCompletionViewModel();
+
+        try
+        {
+            var fileService = _services.GetRequiredService<IBotSharpFileService>();
+            var content = await fileService.InstructPdf(input.Provider, input.Model, input.ModelId, input.Text, input.Files);
+            viewModel.Content = content;
+            return viewModel;
+        }
+        catch (Exception ex)
+        {
+            var error = $"Error in pdf completion. {ex.Message}";
+            _logger.LogError(error);
+            viewModel.Message = error;
+            return viewModel;
+        }
+    }
 }

src/Infrastructure/BotSharp.OpenAPI/ViewModels/Instructs/ImageGenerationViewModel.cs

@@ -4,15 +4,38 @@ namespace BotSharp.OpenAPI.ViewModels.Instructs;
 
 public class ImageGenerationViewModel
 {
-    [JsonPropertyName("revised_prompt")]
-    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
-    public string? RevisedPrompt { get; set; }
+    [JsonPropertyName("content")]
+    public string Content { get; set; } = string.Empty;
 
-    [JsonPropertyName("data")]
+    [JsonPropertyName("images")]
     [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
-    public object? Data { get; set; }
+    public IEnumerable<ImageViewModel> Images { get; set; } = new List<ImageViewModel>();
 
     [JsonPropertyName("message")]
     [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
     public string? Message { get; set; }
 }
+
+public class ImageViewModel
+{
+    [JsonPropertyName("image_url")]
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public string? ImageUrl { get; set; }
+
+    [JsonPropertyName("image_data")]
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public string? ImageData { get; set; }
+
+    [JsonPropertyName("description")]
+    public string Description { get; set; } = string.Empty;
+
+    public static ImageViewModel ToViewModel(ImageGeneration image)
+    {
+        return new ImageViewModel
+        {
+            ImageUrl = image.ImageUrl,
+            ImageData = image.ImageData,
+            Description = image.Description
+        };
+    }
+}

src/Infrastructure/BotSharp.OpenAPI/ViewModels/Instructs/PdfCompletionViewModel.cs

@@ -0,0 +1,13 @@
+using System.Text.Json.Serialization;
+
+namespace BotSharp.OpenAPI.ViewModels.Instructs;
+
+public class PdfCompletionViewModel
+{
+    [JsonPropertyName("content")]
+    public string Content { get; set; } = string.Empty;
+
+    [JsonPropertyName("message")]
+    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
+    public string? Message { get; set; }
+}

src/Plugins/BotSharp.Plugin.AzureOpenAI/Providers/Image/ImageGenerationProvider.cs

@@ -1,3 +1,4 @@
+using BotSharp.Abstraction.Files.Models;
 using OpenAI.Images;
 
 namespace BotSharp.Plugin.AzureOpenAI.Providers.Image;
@@ -8,6 +9,9 @@ public class ImageGenerationProvider : IImageGeneration
     protected readonly IServiceProvider _services;
     protected readonly ILogger _logger;
 
+    private const int DEFAULT_IMAGE_COUNT = 1;
+    private const int IMAGE_COUNT_LIMIT = 3;
+
     protected string _model;
 
     public virtual string Provider => "azure-openai";
@@ -34,23 +38,43 @@ public class ImageGenerationProvider : IImageGeneration
         }
 
         var client = ProviderHelper.GetClient(Provider, _model, _services);
-        var (prompt, options) = PrepareOptions(conversations);
+        var (prompt, imageCount, options) = PrepareOptions(conversations);
         var imageClient = client.GetImageClient(_model);
 
-        var response = imageClient.GenerateImage(prompt, options);
-        var value = response.Value;
+        var response = imageClient.GenerateImages(prompt, imageCount, options);
+        var values = response.Value;
 
         var content = string.Empty;
-        if (!string.IsNullOrEmpty(value.RevisedPrompt))
+        var images = new List<ImageGeneration>();
+        foreach (var value in values)
         {
-            content = value.RevisedPrompt;
+            if (value == null) continue;
+
+            var image = new ImageGeneration { Description = value?.RevisedPrompt ?? string.Empty };
+            if (options.ResponseFormat == GeneratedImageFormat.Uri)
+            {
+                image.ImageUrl = value?.ImageUri?.AbsoluteUri ?? string.Empty;
+            }
+            else if (options.ResponseFormat == GeneratedImageFormat.Bytes)
+            {
+                var base64Str = string.Empty;
+                var bytes = value?.ImageBytes?.ToArray();
+                if (!bytes.IsNullOrEmpty())
+                {
+                    base64Str = Convert.ToBase64String(bytes);
+                }
+                image.ImageData = base64Str;
+            }
+
+            images.Add(image);
+            content += $"{image.Description}\r\n";
         }
 
         var responseMessage = new RoleDialogModel(AgentRole.Assistant, content)
         {
             CurrentAgentId = agent.Id,
             MessageId = conversations.LastOrDefault()?.MessageId ?? string.Empty,
-            Data = options.ResponseFormat == GeneratedImageFormat.Uri ? value.ImageUri?.AbsoluteUri : value.ImageBytes
+            GeneratedImages = images
         };
 
         // After
@@ -69,7 +93,7 @@ public class ImageGenerationProvider : IImageGeneration
         return responseMessage;
     }
 
-    private (string, ImageGenerationOptions) PrepareOptions(List<RoleDialogModel> conversations)
+    private (string, int, ImageGenerationOptions) PrepareOptions(List<RoleDialogModel> conversations)
     {
         var prompt = conversations.LastOrDefault()?.Payload ?? conversations.LastOrDefault()?.Content ?? string.Empty;
 
@@ -77,15 +101,17 @@ public class ImageGenerationProvider : IImageGeneration
         var size = state.GetState("image_size");
         var quality = state.GetState("image_quality");
         var style = state.GetState("image_style");
+        var format = state.GetState("image_format");
+        var count = GetImageCount(state.GetState("image_count", "1"));
 
         var options = new ImageGenerationOptions
         {
             Size = GetImageSize(size),
             Quality = GetImageQuality(quality),
             Style = GetImageStyle(style),
-            ResponseFormat = GeneratedImageFormat.Uri
+            ResponseFormat = GetImageFormat(format)
         };
-        return (prompt, options);
+        return (prompt, count, options);
     }
 
     public void SetModelName(string model)
@@ -164,4 +190,35 @@ public class ImageGenerationProvider : IImageGeneration
 
         return retStyle;
     }
+
+    private GeneratedImageFormat GetImageFormat(string format)
+    {
+        var value = !string.IsNullOrEmpty(format) ? format : "uri";
+
+        GeneratedImageFormat retFormat;
+        switch (value)
+        {
+            case "uri":
+                retFormat = GeneratedImageFormat.Uri;
+                break;
+            case "bytes":
+                retFormat = GeneratedImageFormat.Bytes;
+                break;
+            default:
+                retFormat = GeneratedImageFormat.Uri;
+                break;
+        }
+
+        return retFormat;
+    }
+
+    private int GetImageCount(string count)
+    {
+        if (!int.TryParse(count, out var retCount))
+        {
+            return DEFAULT_IMAGE_COUNT;
+        }
+
+        return retCount > 0 && retCount <= IMAGE_COUNT_LIMIT ? retCount : DEFAULT_IMAGE_COUNT;
+    }
 }