Add otel middleware for IImageGenerator (dotnet#6809)

Author: stephentoub

src/Libraries/Microsoft.Extensions.AI.Abstractions/Image/ImageGenerationOptions.cs

@@ -59,12 +59,16 @@ public class ImageGenerationOptions
     /// </summary>
     public ImageGenerationResponseFormat? ResponseFormat { get; set; }
 
+    /// <summary>Gets or sets any additional properties associated with the options.</summary>
+    public AdditionalPropertiesDictionary? AdditionalProperties { get; set; }
+
     /// <summary>Produces a clone of the current <see cref="ImageGenerationOptions"/> instance.</summary>
     /// <returns>A clone of the current <see cref="ImageGenerationOptions"/> instance.</returns>
     public virtual ImageGenerationOptions Clone()
     {
         ImageGenerationOptions options = new()
         {
+            AdditionalProperties = AdditionalProperties?.Clone(),
             Count = Count,
             MediaType = MediaType,
             ImageSize = ImageSize,

src/Libraries/Microsoft.Extensions.AI.Abstractions/Image/ImageGenerationResponse.cs

@@ -51,4 +51,7 @@ public IList<AIContent> Contents
         get => _contents ??= [];
         set => _contents = value;
     }
+
+    /// <summary>Gets or sets usage details for the image generation response.</summary>
+    public UsageDetails? Usage { get; set; }
 }

src/Libraries/Microsoft.Extensions.AI.OpenAI/OpenAIImageGenerator.cs

@@ -163,9 +163,28 @@ private static ImageGenerationResponse ToImageGenerationResponse(GeneratedImageC
             }
         }
 
+        UsageDetails? ud = null;
+        if (generatedImages.Usage is { } usage)
+        {
+            ud = new()
+            {
+                InputTokenCount = usage.InputTokenCount,
+                OutputTokenCount = usage.OutputTokenCount,
+                TotalTokenCount = usage.TotalTokenCount,
+            };
+
+            if (usage.InputTokenDetails is { } inputDetails)
+            {
+                ud.AdditionalCounts ??= [];
+                ud.AdditionalCounts.Add($"{nameof(usage.InputTokenDetails)}.{nameof(inputDetails.ImageTokenCount)}", inputDetails.ImageTokenCount);
+                ud.AdditionalCounts.Add($"{nameof(usage.InputTokenDetails)}.{nameof(inputDetails.TextTokenCount)}", inputDetails.TextTokenCount);
+            }
+        }
+
         return new ImageGenerationResponse(contents)
         {
-            RawRepresentation = generatedImages
+            RawRepresentation = generatedImages,
+            Usage = ud,
         };
     }
 

src/Libraries/Microsoft.Extensions.AI/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## NOT YET RELEASED
 
 - Updated the EnableSensitiveData properties on OpenTelemetryChatClient/EmbeddingGenerator to respect a OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT environment variable.
+- Added OpenTelemetryImageGenerator to provide OpenTelemetry instrumentation for IImageGenerator implementations.
 
 ## 9.9.0
 

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/OpenTelemetryChatClient.cs

@@ -217,7 +217,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
         }
     }
 
-    private static string SerializeChatMessages(IEnumerable<ChatMessage> messages, ChatFinishReason? chatFinishReason = null)
+    internal static string SerializeChatMessages(IEnumerable<ChatMessage> messages, ChatFinishReason? chatFinishReason = null)
     {
         List<object> output = [];
 
@@ -244,6 +244,12 @@ private static string SerializeChatMessages(IEnumerable<ChatMessage> messages, C
             {
                 switch (content)
                 {
+                    // These are all specified in the convention:
+
+                    case TextContent tc when !string.IsNullOrWhiteSpace(tc.Text):
+                        m.Parts.Add(new OtelGenericPart { Content = tc.Text });
+                        break;
+
                     case FunctionCallContent fcc:
                         m.Parts.Add(new OtelToolCallRequestPart
                         {
@@ -261,8 +267,30 @@ private static string SerializeChatMessages(IEnumerable<ChatMessage> messages, C
                         });
                         break;
 
-                    case TextContent tc:
-                        m.Parts.Add(new OtelGenericPart { Content = tc.Text });
+                    // These are non-standard and are using the "generic" non-text part that provides an extensibility mechanism:
+
+                    case TextReasoningContent trc when !string.IsNullOrWhiteSpace(trc.Text):
+                        m.Parts.Add(new OtelGenericPart { Type = "reasoning", Content = trc.Text });
+                        break;
+
+                    case UriContent uc:
+                        m.Parts.Add(new OtelGenericPart { Type = "image", Content = uc.Uri.ToString() });
+                        break;
+
+                    case DataContent dc:
+                        m.Parts.Add(new OtelGenericPart { Type = "image", Content = dc.Uri });
+                        break;
+
+                    case HostedFileContent fc:
+                        m.Parts.Add(new OtelGenericPart { Type = "file", Content = fc.FileId });
+                        break;
+
+                    case HostedVectorStoreContent vsc:
+                        m.Parts.Add(new OtelGenericPart { Type = "vector_store", Content = vsc.VectorStoreId });
+                        break;
+
+                    case ErrorContent ec:
+                        m.Parts.Add(new OtelGenericPart { Type = "error", Content = ec.Message });
                         break;
 
                     default:
@@ -293,7 +321,7 @@ private static string SerializeChatMessages(IEnumerable<ChatMessage> messages, C
                 string.IsNullOrWhiteSpace(modelId) ? OpenTelemetryConsts.GenAI.Chat : $"{OpenTelemetryConsts.GenAI.Chat} {modelId}",
                 ActivityKind.Client);
 
-            if (activity is not null)
+            if (activity is { IsAllDataRequested: true })
             {
                 _ = activity
                     .AddTag(OpenTelemetryConsts.GenAI.Operation.Name, OpenTelemetryConsts.GenAI.Chat)
@@ -411,15 +439,15 @@ private void TraceResponse(
                 TagList tags = default;
                 tags.Add(OpenTelemetryConsts.GenAI.Token.Type, OpenTelemetryConsts.TokenTypeInput);
                 AddMetricTags(ref tags, requestModelId, response);
-                _tokenUsageHistogram.Record((int)inputTokens);
+                _tokenUsageHistogram.Record((int)inputTokens, tags);
             }
 
             if (usage.OutputTokenCount is long outputTokens)
             {
                 TagList tags = default;
                 tags.Add(OpenTelemetryConsts.GenAI.Token.Type, OpenTelemetryConsts.TokenTypeOutput);
                 AddMetricTags(ref tags, requestModelId, response);
-                _tokenUsageHistogram.Record((int)outputTokens);
+                _tokenUsageHistogram.Record((int)outputTokens, tags);
             }
         }