Improve FunctionInvokingChatClient's awareness of agents (#7030)

This addresses several issues:
1. When FunctionInvokingChatClient is used from an agent, and the agent creates an invoke_agent span, that effectively serves the same purpose as the orchestrate_tools span the FICC would have otherwise created, yielding duplication in the telemetry. If we detect an "invoke_agent" span is the parent, we can simply not create "orchestrate_tools".
2. When FunctionInvokingChatClient is used from an agent, its operation is much more closely tied to the agent than to the inner chat client it's orchestrating. So whereas by default it gets its ActivitySource from the inner client, if it detects an "invoke_agent" span as current, it will instead use that activity's source. This means that we can still get "execute_tool" spans created even if telemetry isn't enabled for the inner chat client but is for the parent agent.
3. "execute_tool" optionally has sensitive data to emit. FICC currently determines whether to emit sensitive data based on the inner OpenTelemetryChatClient, since it's also using its ActivitySource. Now when it uses the "invoke_agent" span, it also picks up the sensitivity setting from a custom property on that span.

Author: stephentoub

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

@@ -15,6 +15,7 @@
 using Microsoft.Shared.Diagnostics;
 
 #pragma warning disable CA2213 // Disposable fields should be disposed
+#pragma warning disable S2219 // Runtime type checking should be simplified
 #pragma warning disable S3353 // Unchanged local variables should be "const"
 
 namespace Microsoft.Extensions.AI;
@@ -268,8 +269,9 @@ public override async Task<ChatResponse> GetResponseAsync(
         _ = Throw.IfNull(messages);
 
         // A single request into this GetResponseAsync may result in multiple requests to the inner client.
-        // Create an activity to group them together for better observability.
-        using Activity? activity = _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName);
+        // Create an activity to group them together for better observability. If there's already a genai "invoke_agent"
+        // span that's current, however, we just consider that the group and don't add a new one.
+        using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName);
 
         // Copy the original messages in order to avoid enumerating the original messages multiple times.
         // The IEnumerable can represent an arbitrary amount of work.
@@ -407,8 +409,9 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
         _ = Throw.IfNull(messages);
 
         // A single request into this GetStreamingResponseAsync may result in multiple requests to the inner client.
-        // Create an activity to group them together for better observability.
-        using Activity? activity = _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName);
+        // Create an activity to group them together for better observability. If there's already a genai "invoke_agent"
+        // span that's current, however, we just consider that the group and don't add a new one.
+        using Activity? activity = CurrentActivityIsInvokeAgent ? null : _activitySource?.StartActivity(OpenTelemetryConsts.GenAI.OrchestrateToolsName);
         UsageDetails? totalUsage = activity is { IsAllDataRequested: true } ? new() : null; // tracked usage across all turns, to be used for activity purposes
 
         // Copy the original messages in order to avoid enumerating the original messages multiple times.
@@ -1108,6 +1111,10 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
         }
     }
 
+    /// <summary>Gets a value indicating whether <see cref="Activity.Current"/> represents an "invoke_agent" span.</summary>
+    private static bool CurrentActivityIsInvokeAgent =>
+        Activity.Current?.DisplayName == OpenTelemetryConsts.GenAI.InvokeAgentName;
+
     /// <summary>Invokes the function asynchronously.</summary>
     /// <param name="context">
     /// The function invocation context detailing the function to be invoked and its arguments along with additional request information.
@@ -1119,7 +1126,12 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
     {
         _ = Throw.IfNull(context);
 
-        using Activity? activity = _activitySource?.StartActivity(
+        // We have multiple possible ActivitySource's we could use. In a chat scenario, we ask the inner client whether it has an ActivitySource.
+        // In an agent scenario, we use the ActivitySource from the surrounding "invoke_agent" activity.
+        Activity? invokeAgentActivity = CurrentActivityIsInvokeAgent ? Activity.Current : null;
+        ActivitySource? source = invokeAgentActivity?.Source ?? _activitySource;
+
+        using Activity? activity = source?.StartActivity(
             $"{OpenTelemetryConsts.GenAI.ExecuteToolName} {context.Function.Name}",
             ActivityKind.Internal,
             default(ActivityContext),
@@ -1133,7 +1145,14 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
 
         long startingTimestamp = Stopwatch.GetTimestamp();
 
-        bool enableSensitiveData = activity is { IsAllDataRequested: true } && InnerClient.GetService<OpenTelemetryChatClient>()?.EnableSensitiveData is true;
+        // If we're in the chat scenario, we determine whether sensitive data is enabled by querying the inner chat client.
+        // If we're in the agent scenario, we determine whether sensitive data is enabled by checking for the relevant custom property on the activity.
+        bool enableSensitiveData =
+            activity is { IsAllDataRequested: true } &&
+            (invokeAgentActivity is not null ?
+             invokeAgentActivity.GetCustomProperty(OpenTelemetryChatClient.SensitiveDataEnabledCustomKey) as string is OpenTelemetryChatClient.SensitiveDataEnabledTrueValue :
+             InnerClient.GetService<OpenTelemetryChatClient>()?.EnableSensitiveData is true);
+
         bool traceLoggingEnabled = _logger.IsEnabled(LogLevel.Trace);
         bool loggedInvoke = false;
         if (enableSensitiveData || traceLoggingEnabled)

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

@@ -30,6 +30,9 @@ namespace Microsoft.Extensions.AI;
 /// </remarks>
 public sealed partial class OpenTelemetryChatClient : DelegatingChatClient
 {
+    internal const string SensitiveDataEnabledCustomKey = "__EnableSensitiveData__";
+    internal const string SensitiveDataEnabledTrueValue = "true";
+
     private readonly ActivitySource _activitySource;
     private readonly Meter _meter;
 
@@ -373,6 +376,11 @@ internal static string SerializeChatMessages(
                 string.IsNullOrWhiteSpace(modelId) ? OpenTelemetryConsts.GenAI.ChatName : $"{OpenTelemetryConsts.GenAI.ChatName} {modelId}",
                 ActivityKind.Client);
 
+            if (EnableSensitiveData)
+            {
+                activity?.SetCustomProperty(SensitiveDataEnabledCustomKey, SensitiveDataEnabledTrueValue);
+            }
+
             if (activity is { IsAllDataRequested: true })
             {
                 _ = activity

src/Libraries/Microsoft.Extensions.AI/OpenTelemetryConsts.cs

@@ -35,6 +35,7 @@ public static class GenAI
         public const string ChatName = "chat";
         public const string EmbeddingsName = "embeddings";
         public const string ExecuteToolName = "execute_tool";
+        public const string InvokeAgentName = "invoke_agent";
         public const string OrchestrateToolsName = "orchestrate_tools"; // Non-standard
         public const string GenerateContentName = "generate_content";
 

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

@@ -15,6 +15,7 @@
 using Xunit;
 
 #pragma warning disable SA1118 // Parameter should not span multiple lines
+#pragma warning disable SA1204 // Static elements should appear before instance elements
 
 namespace Microsoft.Extensions.AI;
 
@@ -1232,6 +1233,220 @@ public async Task ClonesChatOptionsAndResetContinuationTokenForBackgroundRespons
         Assert.Null(actualChatOptions!.ContinuationToken);
     }
 
+    [Fact]
+    public async Task DoesNotCreateOrchestrateToolsSpanWhenInvokeAgentIsParent()
+    {
+        string agentSourceName = Guid.NewGuid().ToString();
+        string clientSourceName = Guid.NewGuid().ToString();
+
+        List<ChatMessage> plan =
+        [
+            new ChatMessage(ChatRole.User, "hello"),
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1")]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", result: "Result 1")]),
+            new ChatMessage(ChatRole.Assistant, "world"),
+        ];
+
+        ChatOptions options = new()
+        {
+            Tools = [AIFunctionFactory.Create(() => "Result 1", "Func1")]
+        };
+
+        Func<ChatClientBuilder, ChatClientBuilder> configure = b => b.Use(c =>
+            new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName)));
+
+        var activities = new List<Activity>();
+
+        using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(agentSourceName)
+            .AddSource(clientSourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        using (var agentSource = new ActivitySource(agentSourceName))
+        using (var invokeAgentActivity = agentSource.StartActivity("invoke_agent"))
+        {
+            Assert.NotNull(invokeAgentActivity);
+            await InvokeAndAssertAsync(options, plan, configurePipeline: configure);
+        }
+
+        Assert.DoesNotContain(activities, a => a.DisplayName == "orchestrate_tools");
+        Assert.Contains(activities, a => a.DisplayName == "chat");
+        Assert.Contains(activities, a => a.DisplayName == "execute_tool Func1");
+
+        var invokeAgent = Assert.Single(activities, a => a.DisplayName == "invoke_agent");
+        var childActivities = activities.Where(a => a != invokeAgent).ToList();
+        Assert.All(childActivities, activity => Assert.Same(invokeAgent, activity.Parent));
+    }
+
+    [Fact]
+    public async Task UsesAgentActivitySourceWhenInvokeAgentIsParent()
+    {
+        string agentSourceName = Guid.NewGuid().ToString();
+        string clientSourceName = Guid.NewGuid().ToString();
+
+        List<ChatMessage> plan =
+        [
+            new ChatMessage(ChatRole.User, "hello"),
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1")]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", result: "Result 1")]),
+            new ChatMessage(ChatRole.Assistant, "world"),
+        ];
+
+        ChatOptions options = new()
+        {
+            Tools = [AIFunctionFactory.Create(() => "Result 1", "Func1")]
+        };
+
+        Func<ChatClientBuilder, ChatClientBuilder> configure = b => b.Use(c =>
+            new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName)));
+
+        var activities = new List<Activity>();
+
+        using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(agentSourceName)
+            .AddSource(clientSourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        using (var agentSource = new ActivitySource(agentSourceName))
+        using (var invokeAgentActivity = agentSource.StartActivity("invoke_agent"))
+        {
+            Assert.NotNull(invokeAgentActivity);
+            await InvokeAndAssertAsync(options, plan, configurePipeline: configure);
+        }
+
+        var executeToolActivities = activities.Where(a => a.DisplayName == "execute_tool Func1").ToList();
+        Assert.NotEmpty(executeToolActivities);
+        Assert.All(executeToolActivities, executeTool => Assert.Equal(agentSourceName, executeTool.Source.Name));
+    }
+
+    public static IEnumerable<object[]> SensitiveDataPropagatesFromAgentActivityWhenInvokeAgentIsParent_MemberData() =>
+        from invokeAgentSensitiveData in new bool?[] { null, false, true }
+        from innerOpenTelemetryChatClient in new bool?[] { null, false, true }
+        select new object?[] { invokeAgentSensitiveData, innerOpenTelemetryChatClient };
+
+    [Theory]
+    [MemberData(nameof(SensitiveDataPropagatesFromAgentActivityWhenInvokeAgentIsParent_MemberData))]
+    public async Task SensitiveDataPropagatesFromAgentActivityWhenInvokeAgentIsParent(
+        bool? invokeAgentSensitiveData, bool? innerOpenTelemetryChatClient)
+    {
+        string agentSourceName = Guid.NewGuid().ToString();
+        string clientSourceName = Guid.NewGuid().ToString();
+
+        List<ChatMessage> plan =
+        [
+            new ChatMessage(ChatRole.User, "hello"),
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1", new Dictionary<string, object?> { ["arg1"] = "secret" })]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", result: "Result 1")]),
+            new ChatMessage(ChatRole.Assistant, "world"),
+        ];
+
+        ChatOptions options = new()
+        {
+            Tools = [AIFunctionFactory.Create(() => "Result 1", "Func1")]
+        };
+
+        var activities = new List<Activity>();
+
+        using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(agentSourceName)
+            .AddSource(clientSourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        using (var agentSource = new ActivitySource(agentSourceName))
+        using (var invokeAgentActivity = agentSource.StartActivity("invoke_agent"))
+        {
+            if (invokeAgentSensitiveData is not null)
+            {
+                invokeAgentActivity?.SetCustomProperty("__EnableSensitiveData__", invokeAgentSensitiveData is true ? "true" : "false");
+            }
+
+            await InvokeAndAssertAsync(options, plan, configurePipeline: b =>
+            {
+                b.UseFunctionInvocation();
+
+                if (innerOpenTelemetryChatClient is not null)
+                {
+                    b.UseOpenTelemetry(sourceName: clientSourceName, configure: c =>
+                    {
+                        c.EnableSensitiveData = innerOpenTelemetryChatClient.Value;
+                    });
+                }
+
+                return b;
+            });
+        }
+
+        var executeToolActivity = Assert.Single(activities, a => a.DisplayName == "execute_tool Func1");
+
+        var hasArguments = executeToolActivity.Tags.Any(t => t.Key == "gen_ai.tool.call.arguments");
+        var hasResult = executeToolActivity.Tags.Any(t => t.Key == "gen_ai.tool.call.result");
+
+        if (invokeAgentSensitiveData is true)
+        {
+            Assert.True(hasArguments, "Expected arguments to be logged when agent EnableSensitiveData is true");
+            Assert.True(hasResult, "Expected result to be logged when agent EnableSensitiveData is true");
+
+            var argsTag = Assert.Single(executeToolActivity.Tags, t => t.Key == "gen_ai.tool.call.arguments");
+            Assert.Contains("arg1", argsTag.Value);
+        }
+        else
+        {
+            Assert.False(hasArguments, "Expected arguments NOT to be logged when agent EnableSensitiveData is false");
+            Assert.False(hasResult, "Expected result NOT to be logged when agent EnableSensitiveData is false");
+        }
+    }
+
+    [Theory]
+    [InlineData(false)]
+    [InlineData(true)]
+    public async Task CreatesOrchestrateToolsSpanWhenNoInvokeAgentParent(bool streaming)
+    {
+        string clientSourceName = Guid.NewGuid().ToString();
+
+        List<ChatMessage> plan =
+        [
+            new ChatMessage(ChatRole.User, "hello"),
+            new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1")]),
+            new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", result: "Result 1")]),
+            new ChatMessage(ChatRole.Assistant, "world"),
+        ];
+
+        ChatOptions options = new()
+        {
+            Tools = [AIFunctionFactory.Create(() => "Result 1", "Func1")]
+        };
+
+        Func<ChatClientBuilder, ChatClientBuilder> configure = b => b.Use(c =>
+            new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName)));
+
+        var activities = new List<Activity>();
+        using TracerProvider tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
+            .AddSource(clientSourceName)
+            .AddInMemoryExporter(activities)
+            .Build();
+
+        if (streaming)
+        {
+            await InvokeAndAssertStreamingAsync(options, plan, configurePipeline: configure);
+        }
+        else
+        {
+            await InvokeAndAssertAsync(options, plan, configurePipeline: configure);
+        }
+
+        var orchestrateTools = Assert.Single(activities, a => a.DisplayName == "orchestrate_tools");
+
+        var executeTools = activities.Where(a => a.DisplayName.StartsWith("execute_tool")).ToList();
+        Assert.NotEmpty(executeTools);
+        foreach (var executeTool in executeTools)
+        {
+            Assert.Same(orchestrateTools, executeTool.Parent);
+        }
+    }
+
     private sealed class CustomSynchronizationContext : SynchronizationContext
     {
         public override void Post(SendOrPostCallback d, object? state)