Add a grouping Activity to FunctionInvokingChatClient (#5777)

* Add a grouping Activity to FunctionInvokingChatClient

* Update test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs

Co-authored-by: Steve Sanderson <[email protected]>

---------

Co-authored-by: Steve Sanderson <[email protected]>

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

@@ -191,6 +191,10 @@ public override async Task<ChatCompletion> CompleteAsync(IList<ChatMessage> chat
     {
         _ = Throw.IfNull(chatMessages);
 
+        // A single request into this CompleteAsync may result in multiple requests to the inner client.
+        // Create an activity to group them together for better observability.
+        using Activity? activity = _activitySource?.StartActivity(nameof(FunctionInvokingChatClient));
+
         ChatCompletion? response = null;
         HashSet<ChatMessage>? messagesToRemove = null;
         HashSet<AIContent>? contentsToRemove = null;
@@ -307,6 +311,10 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
     {
         _ = Throw.IfNull(chatMessages);
 
+        // A single request into this CompleteStreamingAsync may result in multiple requests to the inner client.
+        // Create an activity to group them together for better observability.
+        using Activity? activity = _activitySource?.StartActivity(nameof(FunctionInvokingChatClient));
+
         HashSet<ChatMessage>? messagesToRemove = null;
         List<FunctionCallContent> functionCallContents = [];
         int? choice;
@@ -349,6 +357,7 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
                     }
 
                     yield return update;
+                    Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802
                 }
 
                 // If there are no tools to call, or for any other reason we should stop, return the response.

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

@@ -197,6 +197,7 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
 
                 trackedUpdates.Add(update);
                 yield return update;
+                Activity.Current = activity; // workaround for https://github.com/dotnet/runtime/issues/47802
             }
         }
         finally

test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs

@@ -485,7 +485,14 @@ async Task InvokeAsync(Func<Task> work)
                 Assert.Collection(activities,
                     activity => Assert.Equal("chat", activity.DisplayName),
                     activity => Assert.Equal("Func1", activity.DisplayName),
-                    activity => Assert.Equal("chat", activity.DisplayName));
+                    activity => Assert.Equal("chat", activity.DisplayName),
+                    activity => Assert.Equal(nameof(FunctionInvokingChatClient), activity.DisplayName));
+
+                for (int i = 0; i < activities.Count - 1; i++)
+                {
+                    // Activities are exported in the order of completion, so all except the last are children of the last (i.e., outer)
+                    Assert.Same(activities[activities.Count - 1], activities[i].Parent);
+                }
             }
             else
             {