Skip to content

Commit d7c32f8

Browse files
stephentoubstephentoub
committed
Add logging/activities to FunctionInvokingChatClient (dotnet#5596)
* Add logging/activities to FunctionInvokingChatClient * Change FunctionInvokingChatClient to use ActivitySource from OpenTelemetryChatClient
1 parent ca7d3f2 commit d7c32f8

File tree

9 files changed

+271
-23
lines changed

9 files changed

+271
-23
lines changed

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

Lines changed: 100 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,12 @@
88
using System.Runtime.CompilerServices;
99
using System.Threading;
1010
using System.Threading.Tasks;
11+
using Microsoft.Extensions.Logging;
12+
using Microsoft.Extensions.Logging.Abstractions;
1113
using Microsoft.Shared.Diagnostics;
1214

15+
#pragma warning disable CA2213 // Disposable fields should be disposed
16+
1317
namespace Microsoft.Extensions.AI;
1418

1519
/// <summary>
@@ -34,18 +38,28 @@ namespace Microsoft.Extensions.AI;
3438
/// invocation requests to that same function.
3539
/// </para>
3640
/// </remarks>
37-
public class FunctionInvokingChatClient : DelegatingChatClient
41+
public partial class FunctionInvokingChatClient : DelegatingChatClient
3842
{
43+
/// <summary>The logger to use for logging information about function invocation.</summary>
44+
private readonly ILogger _logger;
45+
46+
/// <summary>The <see cref="ActivitySource"/> to use for telemetry.</summary>
47+
/// <remarks>This component does not own the instance and should not dispose it.</remarks>
48+
private readonly ActivitySource? _activitySource;
49+
3950
/// <summary>Maximum number of roundtrips allowed to the inner client.</summary>
4051
private int? _maximumIterationsPerRequest;
4152

4253
/// <summary>
4354
/// Initializes a new instance of the <see cref="FunctionInvokingChatClient"/> class.
4455
/// </summary>
4556
/// <param name="innerClient">The underlying <see cref="IChatClient"/>, or the next instance in a chain of clients.</param>
46-
public FunctionInvokingChatClient(IChatClient innerClient)
57+
/// <param name="logger">An <see cref="ILogger"/> to use for logging information about function invocation.</param>
58+
public FunctionInvokingChatClient(IChatClient innerClient, ILogger? logger = null)
4759
: base(innerClient)
4860
{
61+
_logger = logger ?? NullLogger.Instance;
62+
_activitySource = innerClient.GetService<ActivitySource>();
4963
}
5064

5165
/// <summary>
@@ -562,13 +576,95 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
562576
/// </param>
563577
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
564578
/// <returns>The result of the function invocation. This may be null if the function invocation returned null.</returns>
565-
protected virtual Task<object?> InvokeFunctionAsync(FunctionInvocationContext context, CancellationToken cancellationToken)
579+
protected virtual async Task<object?> InvokeFunctionAsync(FunctionInvocationContext context, CancellationToken cancellationToken)
566580
{
567581
_ = Throw.IfNull(context);
568582

569-
return context.Function.InvokeAsync(context.CallContent.Arguments, cancellationToken);
583+
using Activity? activity = _activitySource?.StartActivity(context.Function.Metadata.Name);
584+
585+
long startingTimestamp = 0;
586+
if (_logger.IsEnabled(LogLevel.Debug))
587+
{
588+
startingTimestamp = Stopwatch.GetTimestamp();
589+
if (_logger.IsEnabled(LogLevel.Trace))
590+
{
591+
LogInvokingSensitive(context.Function.Metadata.Name, LoggingHelpers.AsJson(context.CallContent.Arguments, context.Function.Metadata.JsonSerializerOptions));
592+
}
593+
else
594+
{
595+
LogInvoking(context.Function.Metadata.Name);
596+
}
597+
}
598+
599+
object? result = null;
600+
try
601+
{
602+
result = await context.Function.InvokeAsync(context.CallContent.Arguments, cancellationToken).ConfigureAwait(false);
603+
}
604+
catch (Exception e)
605+
{
606+
if (activity is not null)
607+
{
608+
_ = activity.SetTag("error.type", e.GetType().FullName)
609+
.SetStatus(ActivityStatusCode.Error, e.Message);
610+
}
611+
612+
if (e is OperationCanceledException)
613+
{
614+
LogInvocationCanceled(context.Function.Metadata.Name);
615+
}
616+
else
617+
{
618+
LogInvocationFailed(context.Function.Metadata.Name, e);
619+
}
620+
621+
throw;
622+
}
623+
finally
624+
{
625+
if (_logger.IsEnabled(LogLevel.Debug))
626+
{
627+
TimeSpan elapsed = GetElapsedTime(startingTimestamp);
628+
629+
if (result is not null && _logger.IsEnabled(LogLevel.Trace))
630+
{
631+
LogInvocationCompletedSensitive(context.Function.Metadata.Name, elapsed, LoggingHelpers.AsJson(result, context.Function.Metadata.JsonSerializerOptions));
632+
}
633+
else
634+
{
635+
LogInvocationCompleted(context.Function.Metadata.Name, elapsed);
636+
}
637+
}
638+
}
639+
640+
return result;
570641
}
571642

643+
private static TimeSpan GetElapsedTime(long startingTimestamp) =>
644+
#if NET
645+
Stopwatch.GetElapsedTime(startingTimestamp);
646+
#else
647+
new((long)((Stopwatch.GetTimestamp() - startingTimestamp) * ((double)TimeSpan.TicksPerSecond / Stopwatch.Frequency)));
648+
#endif
649+
650+
[LoggerMessage(LogLevel.Debug, "Invoking {MethodName}.", SkipEnabledCheck = true)]
651+
private partial void LogInvoking(string methodName);
652+
653+
[LoggerMessage(LogLevel.Trace, "Invoking {MethodName}({Arguments}).", SkipEnabledCheck = true)]
654+
private partial void LogInvokingSensitive(string methodName, string arguments);
655+
656+
[LoggerMessage(LogLevel.Debug, "{MethodName} invocation completed. Duration: {Duration}", SkipEnabledCheck = true)]
657+
private partial void LogInvocationCompleted(string methodName, TimeSpan duration);
658+
659+
[LoggerMessage(LogLevel.Trace, "{MethodName} invocation completed. Duration: {Duration}. Result: {Result}", SkipEnabledCheck = true)]
660+
private partial void LogInvocationCompletedSensitive(string methodName, TimeSpan duration, string result);
661+
662+
[LoggerMessage(LogLevel.Debug, "{MethodName} invocation canceled.")]
663+
private partial void LogInvocationCanceled(string methodName);
664+
665+
[LoggerMessage(LogLevel.Error, "{MethodName} invocation failed.")]
666+
private partial void LogInvocationFailed(string methodName, Exception error);
667+
572668
/// <summary>Provides context for a function invocation.</summary>
573669
public sealed class FunctionInvocationContext
574670
{

‎src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClientBuilderExtensions.cs

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
// The .NET Foundation licenses this file to you under the MIT license.
33

44
using System;
5+
using Microsoft.Extensions.DependencyInjection;
6+
using Microsoft.Extensions.Logging;
57
using Microsoft.Shared.Diagnostics;
68

79
namespace Microsoft.Extensions.AI;
@@ -16,15 +18,21 @@ public static class FunctionInvokingChatClientBuilderExtensions
1618
/// </summary>
1719
/// <remarks>This works by adding an instance of <see cref="FunctionInvokingChatClient"/> with default options.</remarks>
1820
/// <param name="builder">The <see cref="ChatClientBuilder"/> being used to build the chat pipeline.</param>
21+
/// <param name="loggerFactory">An optional <see cref="ILoggerFactory"/> to use to create a logger for logging function invocations.</param>
1922
/// <param name="configure">An optional callback that can be used to configure the <see cref="FunctionInvokingChatClient"/> instance.</param>
2023
/// <returns>The supplied <paramref name="builder"/>.</returns>
21-
public static ChatClientBuilder UseFunctionInvocation(this ChatClientBuilder builder, Action<FunctionInvokingChatClient>? configure = null)
24+
public static ChatClientBuilder UseFunctionInvocation(
25+
this ChatClientBuilder builder,
26+
ILoggerFactory? loggerFactory = null,
27+
Action<FunctionInvokingChatClient>? configure = null)
2228
{
2329
_ = Throw.IfNull(builder);
2430

25-
return builder.Use(innerClient =>
31+
return builder.Use((services, innerClient) =>
2632
{
27-
var chatClient = new FunctionInvokingChatClient(innerClient);
33+
loggerFactory ??= services.GetService<ILoggerFactory>();
34+
35+
var chatClient = new FunctionInvokingChatClient(innerClient, loggerFactory?.CreateLogger(typeof(FunctionInvokingChatClient)));
2836
configure?.Invoke(chatClient);
2937
return chatClient;
3038
});

‎src/Libraries/Microsoft.Extensions.AI/ChatCompletion/LoggingChatClient.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ public override async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteSt
168168
}
169169
}
170170

171-
private string AsJson<T>(T value) => JsonSerializer.Serialize(value, _jsonSerializerOptions.GetTypeInfo(typeof(T)));
171+
private string AsJson<T>(T value) => LoggingHelpers.AsJson(value, _jsonSerializerOptions);
172172

173173
[LoggerMessage(LogLevel.Debug, "{MethodName} invoked.")]
174174
private partial void LogInvoked(string methodName);

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
using Microsoft.Extensions.Logging.Abstractions;
1818
using Microsoft.Shared.Diagnostics;
1919

20+
#pragma warning disable S3358 // Ternary operators should not be nested
21+
2022
namespace Microsoft.Extensions.AI;
2123

2224
/// <summary>A delegating chat client that implements the OpenTelemetry Semantic Conventions for Generative AI systems.</summary>
@@ -106,6 +108,11 @@ protected override void Dispose(bool disposing)
106108
/// </remarks>
107109
public bool EnableSensitiveData { get; set; }
108110

111+
/// <inheritdoc/>
112+
public override object? GetService(Type serviceType, object? serviceKey = null) =>
113+
serviceType == typeof(ActivitySource) ? _activitySource :
114+
base.GetService(serviceType, serviceKey);
115+
109116
/// <inheritdoc/>
110117
public override async Task<ChatCompletion> CompleteAsync(IList<ChatMessage> chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default)
111118
{
@@ -254,7 +261,7 @@ private static ChatCompletion ComposeStreamingUpdatesIntoChatCompletion(
254261
string? modelId = options?.ModelId ?? _modelId;
255262

256263
activity = _activitySource.StartActivity(
257-
$"{OpenTelemetryConsts.GenAI.Chat} {modelId}",
264+
string.IsNullOrWhiteSpace(modelId) ? OpenTelemetryConsts.GenAI.Chat : $"{OpenTelemetryConsts.GenAI.Chat} {modelId}",
258265
ActivityKind.Client);
259266

260267
if (activity is not null)

‎src/Libraries/Microsoft.Extensions.AI/Embeddings/OpenTelemetryEmbeddingGenerator.cs

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -72,13 +72,19 @@ public OpenTelemetryEmbeddingGenerator(IEmbeddingGenerator<TInput, TEmbedding> i
7272
advice: new() { HistogramBucketBoundaries = OpenTelemetryConsts.GenAI.Client.OperationDuration.ExplicitBucketBoundaries });
7373
}
7474

75+
/// <inheritdoc/>
76+
public override object? GetService(Type serviceType, object? serviceKey = null) =>
77+
serviceType == typeof(ActivitySource) ? _activitySource :
78+
base.GetService(serviceType, serviceKey);
79+
7580
/// <inheritdoc/>
7681
public override async Task<GeneratedEmbeddings<TEmbedding>> GenerateAsync(IEnumerable<TInput> values, EmbeddingGenerationOptions? options = null, CancellationToken cancellationToken = default)
7782
{
7883
_ = Throw.IfNull(values);
7984

80-
using Activity? activity = CreateAndConfigureActivity();
85+
using Activity? activity = CreateAndConfigureActivity(options);
8186
Stopwatch? stopwatch = _operationDurationHistogram.Enabled ? Stopwatch.StartNew() : null;
87+
string? requestModelId = options?.ModelId ?? _modelId;
8288

8389
GeneratedEmbeddings<TEmbedding>? response = null;
8490
Exception? error = null;
@@ -93,7 +99,7 @@ public override async Task<GeneratedEmbeddings<TEmbedding>> GenerateAsync(IEnume
9399
}
94100
finally
95101
{
96-
TraceCompletion(activity, response, error, stopwatch);
102+
TraceCompletion(activity, requestModelId, response, error, stopwatch);
97103
}
98104

99105
return response;
@@ -112,18 +118,20 @@ protected override void Dispose(bool disposing)
112118
}
113119

114120
/// <summary>Creates an activity for an embedding generation request, or returns null if not enabled.</summary>
115-
private Activity? CreateAndConfigureActivity()
121+
private Activity? CreateAndConfigureActivity(EmbeddingGenerationOptions? options)
116122
{
117123
Activity? activity = null;
118124
if (_activitySource.HasListeners())
119125
{
126+
string? modelId = options?.ModelId ?? _modelId;
127+
120128
activity = _activitySource.StartActivity(
121-
$"{OpenTelemetryConsts.GenAI.Embed} {_modelId}",
129+
string.IsNullOrWhiteSpace(modelId) ? OpenTelemetryConsts.GenAI.Embed : $"{OpenTelemetryConsts.GenAI.Embed} {modelId}",
122130
ActivityKind.Client,
123131
default(ActivityContext),
124132
[
125133
new(OpenTelemetryConsts.GenAI.Operation.Name, OpenTelemetryConsts.GenAI.Embed),
126-
new(OpenTelemetryConsts.GenAI.Request.Model, _modelId),
134+
new(OpenTelemetryConsts.GenAI.Request.Model, modelId),
127135
new(OpenTelemetryConsts.GenAI.SystemName, _modelProvider),
128136
]);
129137

@@ -149,6 +157,7 @@ protected override void Dispose(bool disposing)
149157
/// <summary>Adds embedding generation response information to the activity.</summary>
150158
private void TraceCompletion(
151159
Activity? activity,
160+
string? requestModelId,
152161
GeneratedEmbeddings<TEmbedding>? embeddings,
153162
Exception? error,
154163
Stopwatch? stopwatch)
@@ -167,7 +176,7 @@ private void TraceCompletion(
167176
if (_operationDurationHistogram.Enabled && stopwatch is not null)
168177
{
169178
TagList tags = default;
170-
AddMetricTags(ref tags, responseModelId);
179+
AddMetricTags(ref tags, requestModelId, responseModelId);
171180
if (error is not null)
172181
{
173182
tags.Add(OpenTelemetryConsts.Error.Type, error.GetType().FullName);
@@ -180,7 +189,7 @@ private void TraceCompletion(
180189
{
181190
TagList tags = default;
182191
tags.Add(OpenTelemetryConsts.GenAI.Token.Type, "input");
183-
AddMetricTags(ref tags, responseModelId);
192+
AddMetricTags(ref tags, requestModelId, responseModelId);
184193

185194
_tokenUsageHistogram.Record(inputTokens.Value);
186195
}
@@ -206,13 +215,13 @@ private void TraceCompletion(
206215
}
207216
}
208217

209-
private void AddMetricTags(ref TagList tags, string? responseModelId)
218+
private void AddMetricTags(ref TagList tags, string? requestModelId, string? responseModelId)
210219
{
211220
tags.Add(OpenTelemetryConsts.GenAI.Operation.Name, OpenTelemetryConsts.GenAI.Embed);
212221

213-
if (_modelId is string requestModel)
222+
if (requestModelId is not null)
214223
{
215-
tags.Add(OpenTelemetryConsts.GenAI.Request.Model, requestModel);
224+
tags.Add(OpenTelemetryConsts.GenAI.Request.Model, requestModelId);
216225
}
217226

218227
tags.Add(OpenTelemetryConsts.GenAI.SystemName, _modelProvider);

‎src/Libraries/Microsoft.Extensions.AI/LoggingHelpers.cs

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
4+
#pragma warning disable CA1031 // Do not catch general exception types
5+
#pragma warning disable S108 // Nested blocks of code should not be left empty
6+
#pragma warning disable S2486 // Generic exceptions should not be ignored
7+
8+
using System.Text.Json;
9+
10+
namespace Microsoft.Extensions.AI;
11+
12+
/// <summary>Provides internal helpers for implementing logging.</summary>
13+
internal static class LoggingHelpers
14+
{
15+
/// <summary>Serializes <paramref name="value"/> as JSON for logging purposes.</summary>
16+
public static string AsJson<T>(T value, JsonSerializerOptions? options)
17+
{
18+
if (options?.TryGetTypeInfo(typeof(T), out var typeInfo) is true ||
19+
AIJsonUtilities.DefaultOptions.TryGetTypeInfo(typeof(T), out typeInfo))
20+
{
21+
try
22+
{
23+
return JsonSerializer.Serialize(value, typeInfo);
24+
}
25+
catch
26+
{
27+
}
28+
}
29+
30+
// If we're unable to get a type info for the value, or if we fail to serialize,
31+
// return an empty JSON object. We do not want lack of type info to disrupt application behavior with exceptions.
32+
return "{}";
33+
}
34+
}

‎test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/TestChatClient.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ public sealed class TestChatClient : IChatClient
1818

1919
public Func<IList<ChatMessage>, ChatOptions?, CancellationToken, IAsyncEnumerable<StreamingChatCompletionUpdate>>? CompleteStreamingAsyncCallback { get; set; }
2020

21-
public Func<Type, object?, object?>? GetServiceCallback { get; set; }
21+
public Func<Type, object?, object?> GetServiceCallback { get; set; } = (_, _) => null;
2222

2323
public Task<ChatCompletion> CompleteAsync(IList<ChatMessage> chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default)
2424
=> CompleteAsyncCallback!.Invoke(chatMessages, options, cancellationToken);
@@ -27,7 +27,7 @@ public IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteStreamingAsync(IL
2727
=> CompleteStreamingAsyncCallback!.Invoke(chatMessages, options, cancellationToken);
2828

2929
public object? GetService(Type serviceType, object? serviceKey = null)
30-
=> GetServiceCallback!(serviceType, serviceKey);
30+
=> GetServiceCallback(serviceType, serviceKey);
3131

3232
void IDisposable.Dispose()
3333
{

‎test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/TestEmbeddingGenerator.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,13 @@ public sealed class TestEmbeddingGenerator : IEmbeddingGenerator<string, Embeddi
1414

1515
public Func<IEnumerable<string>, EmbeddingGenerationOptions?, CancellationToken, Task<GeneratedEmbeddings<Embedding<float>>>>? GenerateAsyncCallback { get; set; }
1616

17-
public Func<Type, object?, object?>? GetServiceCallback { get; set; }
17+
public Func<Type, object?, object?> GetServiceCallback { get; set; } = (_, _) => null;
1818

1919
public Task<GeneratedEmbeddings<Embedding<float>>> GenerateAsync(IEnumerable<string> values, EmbeddingGenerationOptions? options = null, CancellationToken cancellationToken = default)
2020
=> GenerateAsyncCallback!.Invoke(values, options, cancellationToken);
2121

2222
public object? GetService(Type serviceType, object? serviceKey = null)
23-
=> GetServiceCallback!(serviceType, serviceKey);
23+
=> GetServiceCallback(serviceType, serviceKey);
2424

2525
void IDisposable.Dispose()
2626
{

0 commit comments

Comments
 (0)