Skip to content

Commit b234f6f

Browse files
stephentoubstephentoub
committed
Change FunctionInvokingChatClient to use ActivitySource from OpenTelemetryChatClient
1 parent e7a1d5e commit b234f6f

File tree

6 files changed

+51
-38
lines changed

6 files changed

+51
-38
lines changed

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

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212
using Microsoft.Extensions.Logging.Abstractions;
1313
using Microsoft.Shared.Diagnostics;
1414

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

1719
/// <summary>
@@ -38,12 +40,13 @@ namespace Microsoft.Extensions.AI;
3840
/// </remarks>
3941
public partial class FunctionInvokingChatClient : DelegatingChatClient
4042
{
41-
/// <summary><see cref="ActivitySource"/> used for tracking function invocations.</summary>
42-
private static readonly ActivitySource _activitySource = new(typeof(FunctionInvokingChatClient).FullName!);
43-
4443
/// <summary>The logger to use for logging information about function invocation.</summary>
4544
private readonly ILogger _logger;
4645

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+
4750
/// <summary>Maximum number of roundtrips allowed to the inner client.</summary>
4851
private int? _maximumIterationsPerRequest;
4952

@@ -56,6 +59,7 @@ public FunctionInvokingChatClient(IChatClient innerClient, ILogger? logger = nul
5659
: base(innerClient)
5760
{
5861
_logger = logger ?? NullLogger.Instance;
62+
_activitySource = innerClient.GetService<ActivitySource>();
5963
}
6064

6165
/// <summary>
@@ -178,15 +182,6 @@ public int? MaximumIterationsPerRequest
178182
}
179183
}
180184

181-
/// <summary>
182-
/// Gets or sets a value indicating whether <see cref="Activity"/>s should be used to
183-
/// provide telemetry for function invocation.
184-
/// </summary>
185-
/// <remarks>
186-
/// The default value is <see langword="true"/>.
187-
/// </remarks>
188-
public bool EnableTelemetry { get; set; } = true;
189-
190185
/// <inheritdoc/>
191186
public override async Task<ChatCompletion> CompleteAsync(IList<ChatMessage> chatMessages, ChatOptions? options = null, CancellationToken cancellationToken = default)
192187
{
@@ -585,7 +580,7 @@ FunctionResultContent CreateFunctionResultContent(FunctionInvocationResult resul
585580
{
586581
_ = Throw.IfNull(context);
587582

588-
using Activity? activity = EnableTelemetry ? _activitySource.StartActivity(context.Function.Metadata.Name) : null;
583+
using Activity? activity = _activitySource?.StartActivity(context.Function.Metadata.Name);
589584

590585
long startingTimestamp = 0;
591586
if (_logger.IsEnabled(LogLevel.Debug))

‎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);

‎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
{

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

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -352,11 +352,14 @@ await InvokeAndAssertAsync(options, [
352352
[InlineData(true)]
353353
public async Task FunctionInvocationTrackedWithActivity(bool enableTelemetry)
354354
{
355+
string sourceName = Guid.NewGuid().ToString();
355356
var activities = new List<Activity>();
356-
using var tracerProvider = OpenTelemetry.Sdk.CreateTracerProviderBuilder()
357-
.AddSource("Microsoft.Extensions.AI.*")
357+
using TracerProvider? tracerProvider = enableTelemetry ?
358+
OpenTelemetry.Sdk.CreateTracerProviderBuilder()
359+
.AddSource(sourceName)
358360
.AddInMemoryExporter(activities)
359-
.Build();
361+
.Build() :
362+
null;
360363

361364
var options = new ChatOptions
362365
{
@@ -369,16 +372,15 @@ await InvokeAndAssertAsync(options, [
369372
new ChatMessage(ChatRole.Tool, [new FunctionResultContent("callId1", "Func1", result: "Result 1")]),
370373
new ChatMessage(ChatRole.Assistant, "world"),
371374
], configurePipeline: b => b.Use(c =>
372-
new FunctionInvokingChatClient(c) { EnableTelemetry = enableTelemetry }));
375+
new FunctionInvokingChatClient(
376+
new OpenTelemetryChatClient(c, sourceName: sourceName))));
373377

374378
if (enableTelemetry)
375379
{
376-
var activity = Assert.Single(activities);
377-
378-
Assert.NotNull(activity.Id);
379-
Assert.NotEmpty(activity.Id);
380-
381-
Assert.Equal("Func1", activity.DisplayName);
380+
Assert.Collection(activities,
381+
activity => Assert.Equal("chat", activity.DisplayName),
382+
activity => Assert.Equal("Func1", activity.DisplayName),
383+
activity => Assert.Equal("chat", activity.DisplayName));
382384
}
383385
else
384386
{

0 commit comments

Comments
 (0)