From 4247b03a0a7d33c4462e6043eb723bab8bf7e1d5 Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Sun, 9 Nov 2025 14:16:21 +0000 Subject: [PATCH] Improve FunctionInvokingChatClient's awareness of agents 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. --- .../FunctionInvokingChatClient.cs | 31 ++- .../ChatCompletion/OpenTelemetryChatClient.cs | 8 + .../OpenTelemetryConsts.cs | 1 + .../FunctionInvokingChatClientTests.cs | 215 ++++++++++++++++++ 4 files changed, 249 insertions(+), 6 deletions(-) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs index fb821c984df..b8a09e59e4f 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs +++ b/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 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 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 } } + /// Gets a value indicating whether represents an "invoke_agent" span. + private static bool CurrentActivityIsInvokeAgent => + Activity.Current?.DisplayName == OpenTelemetryConsts.GenAI.InvokeAgentName; + /// Invokes the function asynchronously. /// /// 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()?.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()?.EnableSensitiveData is true); + bool traceLoggingEnabled = _logger.IsEnabled(LogLevel.Trace); bool loggedInvoke = false; if (enableSensitiveData || traceLoggingEnabled) diff --git a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/OpenTelemetryChatClient.cs b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/OpenTelemetryChatClient.cs index 1f630a5a62c..264832508ae 100644 --- a/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/OpenTelemetryChatClient.cs +++ b/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/OpenTelemetryChatClient.cs @@ -30,6 +30,9 @@ namespace Microsoft.Extensions.AI; /// 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 diff --git a/src/Libraries/Microsoft.Extensions.AI/OpenTelemetryConsts.cs b/src/Libraries/Microsoft.Extensions.AI/OpenTelemetryConsts.cs index 3def478d0e1..fa7e36333ce 100644 --- a/src/Libraries/Microsoft.Extensions.AI/OpenTelemetryConsts.cs +++ b/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"; diff --git a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs b/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs index 2308a921ab3..4d086ebf61e 100644 --- a/test/Libraries/Microsoft.Extensions.AI.Tests/ChatCompletion/FunctionInvokingChatClientTests.cs +++ b/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 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 configure = b => b.Use(c => + new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName))); + + var activities = new List(); + + 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 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 configure = b => b.Use(c => + new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName))); + + var activities = new List(); + + 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 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 plan = + [ + new ChatMessage(ChatRole.User, "hello"), + new ChatMessage(ChatRole.Assistant, [new FunctionCallContent("callId1", "Func1", new Dictionary { ["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(); + + 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 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 configure = b => b.Use(c => + new FunctionInvokingChatClient(new OpenTelemetryChatClient(c, sourceName: clientSourceName))); + + var activities = new List(); + 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)