Improve content type detection from responses

src/xAI.Tests/ImageGeneratorTests.cs

@@ -74,7 +74,7 @@ public async Task GenerateImage_WithEditsToPreviousImage()
     public async Task GenerateImage_WithBase64Response_ReturnsDataContent()
     {
         var imageGenerator = new GrokClient(Configuration["XAI_API_KEY"]!)
-            .AsIImageGenerator("grok-2-image");
+            .AsIImageGenerator("grok-imagine-image-beta");
 
         var request = new ImageGenerationRequest("A sunset over mountains");
         var options = new ImageGenerationOptions
@@ -99,6 +99,22 @@ public async Task GenerateImage_WithBase64Response_ReturnsDataContent()
         output.WriteLine($"Generated image size: {dataContent.Data.Length} bytes");
     }
 
+    [SecretsFact("XAI_API_KEY")]
+    public async Task GenerateImage_DefaultsToUriContent()
+    {
+        var imageGenerator = new GrokClient(Configuration["XAI_API_KEY"]!)
+            .AsIImageGenerator("grok-imagine-image-beta");
+
+        var request = new ImageGenerationRequest("A sunset over mountains");
+        var response = await imageGenerator.GenerateAsync(request);
+
+        Assert.NotNull(response);
+        Assert.NotEmpty(response.Contents);
+        Assert.Single(response.Contents);
+
+        Assert.IsType<UriContent>(response.Contents.First());
+    }
+
     [SecretsFact("XAI_API_KEY")]
     public async Task GenerateMultipleImages_ReturnsCorrectCount()
     {

src/xAI/GrokImageGenerator.cs

@@ -10,6 +10,18 @@ namespace xAI;
 /// </summary>
 sealed class GrokImageGenerator : IImageGenerator
 {
+    // add inverted dictionary for extension to mime type if needed in future
+    static readonly Dictionary<string, string> extensionToMimeType = new(StringComparer.OrdinalIgnoreCase)
+    {
+        [".png"] = "image/png",
+        [".jpg"] = "image/jpeg",
+        [".jpeg"] = "image/jpeg",
+        [".webp"] = "image/webp",
+        [".gif"] = "image/gif",
+        [".bmp"] = "image/bmp",
+        [".tiff"] = "image/tiff",
+    };
+
     const string DefaultInputContentType = "image/png";
     const string DefaultOutputContentType = "image/jpeg";
 
@@ -50,15 +62,12 @@ public async Task<ImageGenerationResponse> GenerateAsync(
         if (options?.Count is { } count)
             protocolRequest.N = count;
 
-        if (options?.ResponseFormat is { } responseFormat)
+        protocolRequest.Format = (options?.ResponseFormat ?? ImageGenerationResponseFormat.Uri) switch
         {
-            protocolRequest.Format = responseFormat switch
-            {
-                ImageGenerationResponseFormat.Uri => ImageFormat.ImgFormatUrl,
-                ImageGenerationResponseFormat.Data => ImageFormat.ImgFormatBase64,
-                _ => throw new ArgumentException($"Unsupported response format: {responseFormat}", nameof(options))
-            };
-        }
+            ImageGenerationResponseFormat.Uri => ImageFormat.ImgFormatUrl,
+            ImageGenerationResponseFormat.Data => ImageFormat.ImgFormatBase64,
+            _ => throw new ArgumentException($"Unsupported response format: {options?.ResponseFormat}", nameof(options))
+        };
 
         // Handle image editing if original images are provided
         if (request.OriginalImages?.FirstOrDefault() is { } originalImage)
@@ -108,20 +117,25 @@ void IDisposable.Dispose() { }
     static ImageGenerationResponse ToImageGenerationResponse(ImageResponse response)
     {
         var contents = new List<AIContent>();
+        var contentType = DefaultOutputContentType;
 
         foreach (var image in response.Images)
         {
             switch (image.ImageCase)
             {
                 case GeneratedImage.ImageOneofCase.Base64:
                     {
+                        // We assume JPEG since there's no way to get the actual content type.
                         var imageBytes = Convert.FromBase64String(image.Base64);
-                        contents.Add(new DataContent(imageBytes, DefaultOutputContentType));
+                        contents.Add(new DataContent(imageBytes, contentType));
                         break;
                     }
                 case GeneratedImage.ImageOneofCase.Url:
                     {
-                        contents.Add(new UriContent(new Uri(image.Url), DefaultOutputContentType));
+                        if (Path.GetExtension(image.Url) is { } extension && extensionToMimeType.TryGetValue(extension, out var mimeType))
+                            contentType = mimeType;
+
+                        contents.Add(new UriContent(new Uri(image.Url), contentType));
                         break;
                     }
                 default: