[API Proposal]: Trace enrichment for OpenTelemetryChatClient

[jdaigle]

Background and motivation

Support for extending and enriching Activities generated by OpenTelemetryChatClient.

Enrichment is a common pattern for OTEL instrumentation libraries.

API Proposal

// TODO

API Usage

// TODO

Alternative Designs

No response

Risks

No response

[stephentoub]

What did you have in mind?

You can already add a component to the pipeline after the OpenTelemetryChatClient that is then able to augment the span, e.g.

IChatClient client = new OpenAIClient(Environment.GetEnvironmentVariable("AI:OpenAI:ApiKey")).GetOpenAIResponseClient("gpt-4o-mini").AsIChatClient();

client = client
    .AsBuilder()
    .UseOpenTelemetry(sourceName: "MyConsoleApp", configure: c => c.EnableSensitiveData = true)
    .Use(async (messages, options, next, cancellationToken) =>
    {
        try
        {
            await next(messages, options, cancellationToken);
        }
        finally
        {
            // Enrich the OpenTelemetry span with a custom tag
            Activity.Current?.SetTag("my_custom_tag", "my_custom_value");
        }
    })
    .Build(host.Services);

Console.WriteLine(await client.GetResponseAsync("Hello, world!"));

[jdaigle]

@stephentoub I'm particularly interested in adding tags based on AdditionalProperties, which requires access to the ChatResponse. I guess I could do something similar the getResponseFunc and getStreamingResponseFunc callbacks?

Implementing getResponseFunc is pretty easy. But for getStreamingResponseFunc, I want to make sure this will work and won't cause other problems:

public static async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
    IEnumerable<ChatMessage> messages,
    ChatOptions? options,
    IChatClient chatClient,
    [EnumeratorCancellation] CancellationToken cancellationToken)
{
    var innerChangeResponse = chatClient.GetStreamingResponseAsync(messages, options, cancellationToken);
    var trackedUpdates = new List<ChatResponseUpdate>();
    await foreach (var update in innerChangeResponse)
    {
        trackedUpdates.Add(update);
        yield return update;
    }

    var chatResponse = trackedUpdates.ToChatResponse();
    Activity.Current?.SetTag("my_custom_tag", "my_custom_value");
}

Since OpenTelemetryChatClient was already basically doing this - tracking those changes, calling ToChatResponse(), and then updating the Activity based on the response) - I didn't necessarily want to do the same work again. And the implementation of OpenTelemetryChatClient.GetStreamingResponseAsync seems rather complex, and I wasn't sure how much of that complexity is needed.

[stephentoub]

@jdaigle, thanks.

I'm particularly interested in adding tags based on AdditionalProperties

In case it's what you need, note that everything in response.AdditionalProperties is already added as a tag onto the span, if EnableSensitiveData is set to true: https://github.com/dotnet/extensions/blob/c0958881c902f3c8055005dce8e41b627bb600ff/src/Libraries/Microsoft.Extensions.AI/ChatCompletion/OpenTelemetryChatClient.cs#L535-L541

I guess I could do something similar the getResponseFunc and getStreamingResponseFunc callbacks?

Yes, if you need to do something custom based on looking at the response, you could either use those callbacks, or just derive from DelegatingChatClient and override the Get{Streaming}ResponseAsync methods and add that to the pipeline; the Use methods with those delegates is just short-hand for that.

seems rather complex

We could possibly do more, of course.

@JamesNK, @halter73, @tarekgh, what patterns would you recommend we apply here, for enabling developers to augment the spans created by OpenTelemetryChatClient and friends?

[tarekgh]

CC @noahfalk

Users can use OpenTelemetry SDK to enrich any Activity. Just subclass BaseProcessor<Activity> and implement OnStart/OnEnd then register that to the OpenTelemetry pipelines.

using var tracerProvider = Sdk.CreateTracerProviderBuilder()
                .AddSource("MyCompany.MyService")
                .AddProcessor(new CustomActivityEnricher()) // 👈 custom enrichment processor
                .Build();

@jdaigle Is it ok in your scenario to enrich the activity only when EnableSensitiveData is enabled? I am asking because as @stephentoub pointed out the additional properties already added to activity when EnableSensitiveData is enabled and if you use OpenTelemetry OnEnd processor callback, you can check the existing tags and add more accordingly.

[evgenyfedorov2]

see also this https://github.com/open-telemetry/opentelemetry-dotnet-contrib/tree/main/src/OpenTelemetry.Extensions.Enrichment which already contains all BaseProcessor infra, you only need to implement the abstract class TraceEnricher and register it in DI.