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)