[.NET] Parse data uri when creating ImageMessage (#4272)

* add test * parse data uri from ImageMessage
microsoft · Nov 19, 2024 · ac8fd99 · ac8fd99
1 parent 01aed82
commit ac8fd99
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 20 deletions.
diff --git a/dotnet/src/AutoGen.Core/Message/ImageMessage.cs b/dotnet/src/AutoGen.Core/Message/ImageMessage.cs
@@ -2,39 +2,65 @@
 // ImageMessage.cs
 
 using System;
+using System.Text.RegularExpressions;
 
 namespace AutoGen.Core;
 
 public class ImageMessage : IMessage
 {
-    public ImageMessage(Role role, string url, string? from = null, string? mimeType = null)
-        : this(role, new Uri(url), from, mimeType)
-    {
-    }
+    private static readonly Regex s_DataUriRegex = new Regex(@"^data:(?<mediatype>[^;]+);base64,(?<data>.*)$", RegexOptions.Compiled);
 
-    public ImageMessage(Role role, Uri uri, string? from = null, string? mimeType = null)
+    /// <summary>
+    /// Create an ImageMessage from a url.
+    /// The url can be a regular url or a data uri.
+    /// If the url is a data uri, the scheme must be "data" and the format must be data:[<mediatype>][;base64],<data>
+    /// </summary>
+    public ImageMessage(Role role, string url, string? from = null, string? mimeType = null)
     {
         this.Role = role;
         this.From = from;
-        this.Url = uri.ToString();
 
-        // try infer mimeType from uri extension if not provided
-        if (mimeType is null)
+        // url might be a data uri or a regular url
+        if (url.StartsWith("data:", StringComparison.OrdinalIgnoreCase))
+        {
+            // the url must be in the format of data:[<mediatype>][;base64],<data>
+            var match = s_DataUriRegex.Match(url);
+
+            if (!match.Success)
+            {
+                throw new ArgumentException("Invalid DataUri format, expected data:[<mediatype>][;base64],<data>", nameof(url));
+            }
+
+            this.Data = new BinaryData(Convert.FromBase64String(match.Groups["data"].Value), match.Groups["mediatype"].Value);
+
+            this.MimeType = match.Groups["mediatype"].Value;
+        }
+        else
         {
-            mimeType = uri switch
+            this.Url = url;
+            // try infer mimeType from uri extension if not provided
+            if (mimeType is null)
             {
-                _ when uri.AbsoluteUri.EndsWith(".png", StringComparison.OrdinalIgnoreCase) => "image/png",
-                _ when uri.AbsoluteUri.EndsWith(".jpg", StringComparison.OrdinalIgnoreCase) => "image/jpeg",
-                _ when uri.AbsoluteUri.EndsWith(".jpeg", StringComparison.OrdinalIgnoreCase) => "image/jpeg",
-                _ when uri.AbsoluteUri.EndsWith(".gif", StringComparison.OrdinalIgnoreCase) => "image/gif",
-                _ when uri.AbsoluteUri.EndsWith(".bmp", StringComparison.OrdinalIgnoreCase) => "image/bmp",
-                _ when uri.AbsoluteUri.EndsWith(".webp", StringComparison.OrdinalIgnoreCase) => "image/webp",
-                _ when uri.AbsoluteUri.EndsWith(".svg", StringComparison.OrdinalIgnoreCase) => "image/svg+xml",
-                _ => throw new ArgumentException("MimeType is required for ImageMessage", nameof(mimeType))
-            };
+                mimeType = url switch
+                {
+                    _ when url.EndsWith(".png", StringComparison.OrdinalIgnoreCase) => "image/png",
+                    _ when url.EndsWith(".jpg", StringComparison.OrdinalIgnoreCase) => "image/jpeg",
+                    _ when url.EndsWith(".jpeg", StringComparison.OrdinalIgnoreCase) => "image/jpeg",
+                    _ when url.EndsWith(".gif", StringComparison.OrdinalIgnoreCase) => "image/gif",
+                    _ when url.EndsWith(".bmp", StringComparison.OrdinalIgnoreCase) => "image/bmp",
+                    _ when url.EndsWith(".webp", StringComparison.OrdinalIgnoreCase) => "image/webp",
+                    _ when url.EndsWith(".svg", StringComparison.OrdinalIgnoreCase) => "image/svg+xml",
+                    _ => throw new ArgumentException("MimeType is required for ImageMessage", nameof(mimeType))
+                };
+            }
+
+            this.MimeType = mimeType;
         }
+    }
 
-        this.MimeType = mimeType;
+    public ImageMessage(Role role, Uri uri, string? from = null, string? mimeType = null)
+        : this(role, uri.ToString(), from, mimeType)
+    {
     }
 
     public ImageMessage(Role role, BinaryData data, string? from = null)

diff --git a/dotnet/src/AutoGen.SemanticKernel/Middleware/SemanticKernelChatMessageContentConnector.cs b/dotnet/src/AutoGen.SemanticKernel/Middleware/SemanticKernelChatMessageContentConnector.cs
@@ -181,7 +181,19 @@ private IEnumerable<ChatMessageContent> ProcessMessageForOthers(TextMessage mess
     private IEnumerable<ChatMessageContent> ProcessMessageForOthers(ImageMessage message)
     {
         var collectionItems = new ChatMessageContentItemCollection();
-        collectionItems.Add(new ImageContent(new Uri(message.Url ?? message.BuildDataUri())));
+        if (message.Url is not null)
+        {
+            collectionItems.Add(new ImageContent(new Uri(message.Url)));
+        }
+        else if (message.BuildDataUri() is string dataUri)
+        {
+            collectionItems.Add(new ImageContent(dataUri));
+        }
+        else
+        {
+            throw new InvalidOperationException("ImageMessage must have Url or DataUri");
+        }
+
         return [new ChatMessageContent(AuthorRole.User, collectionItems)];
     }
 

diff --git a/dotnet/test/AutoGen.Tests/ImageMessageTests.cs b/dotnet/test/AutoGen.Tests/ImageMessageTests.cs
@@ -35,4 +35,18 @@ public async Task ItCreateFromUrl()
         imageMessage.MimeType.Should().Be("image/png");
         imageMessage.Data.Should().BeNull();
     }
+
+    [Fact]
+    public async Task ItCreateFromBase64Url()
+    {
+        var image = Path.Combine("testData", "images", "background.png");
+        var binary = File.ReadAllBytes(image);
+        var base64 = Convert.ToBase64String(binary);
+
+        var base64Url = $"data:image/png;base64,{base64}";
+        var imageMessage = new ImageMessage(Role.User, base64Url);
+
+        imageMessage.BuildDataUri().Should().Be(base64Url);
+        imageMessage.MimeType.Should().Be("image/png");
+    }
 }