Skip to content

add additional attributes to the OpenAI instrumentation span #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
realark merged 1 commit into from
Dec 19, 2025
Merged

add additional attributes to the OpenAI instrumentation span #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 34 additions & 3 deletions src/Braintrust.Sdk/Instrumentation/OpenAI/InstrumentedOpenAIClient.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,16 +147,20 @@ public override ClientResult<ChatCompletion> CompleteChat(IEnumerable<ChatMessag
{
// Start a span for the chat completion
var activity = _activitySource.StartActivity("Chat Completion", ActivityKind.Client);
var startTime = DateTime.UtcNow;

try
{
// Call the underlying client - this will trigger HTTP call and capture JSON
var result = _client.CompleteChat(messages, options, cancellationToken);

// Calculate time to first token
var timeToFirstToken = (DateTime.UtcNow - startTime).TotalSeconds;

// Get the captured HTTP JSON from Activity baggage
if (activity != null && _captureMessageContent)
{
TagActivity(activity);
TagActivity(activity, timeToFirstToken);
}

return result;
Expand Down Expand Up @@ -186,16 +190,20 @@ public override async Task<ClientResult<ChatCompletion>> CompleteChatAsync(IEnum
{
// Start a span for the chat completion
var activity = _activitySource.StartActivity("Chat Completion", ActivityKind.Client);
var startTime = DateTime.UtcNow;

try
{
// Call the underlying client - this will trigger HTTP call and capture JSON
var result = await _client.CompleteChatAsync(messages, options, cancellationToken).ConfigureAwait(false);

// Calculate time to first token
var timeToFirstToken = (DateTime.UtcNow - startTime).TotalSeconds;

// Get the captured HTTP JSON from Activity baggage
if (activity != null && _captureMessageContent)
{
TagActivity(activity);
TagActivity(activity, timeToFirstToken);
}

return result;
Expand All @@ -219,7 +227,7 @@ public override async Task<ClientResult<ChatCompletion>> CompleteChatAsync(IEnum
}

// TODO: Override other methods as needed (CompleteChatStreaming, etc.)
private void TagActivity(Activity activity)
private void TagActivity(Activity activity, double? timeToFirstToken = null)
{
activity.SetTag("provider", "openai");
{
Expand All @@ -238,6 +246,29 @@ private void TagActivity(Activity activity)
var responseJson = JsonNode.Parse(responseRaw);
activity.SetTag("gen_ai.response.model", responseJson?["model"]?.ToString());
activity.SetTag("braintrust.output_json", responseJson?["choices"]?.ToString());

// Extract token usage metrics
var usage = responseJson?["usage"];
if (usage != null)
{
var promptTokens = usage["prompt_tokens"]?.GetValue<int?>();
var completionTokens = usage["completion_tokens"]?.GetValue<int?>();
var totalTokens = usage["total_tokens"]?.GetValue<int?>();

if (promptTokens.HasValue)
activity.SetTag("braintrust.metrics.prompt_tokens", promptTokens.Value);
if (completionTokens.HasValue)
activity.SetTag("braintrust.metrics.completion_tokens", completionTokens.Value);
if (totalTokens.HasValue)
activity.SetTag("braintrust.metrics.tokens", totalTokens.Value);
}

// Set time_to_first_token metric
// For non-streaming responses, this is the total response time
if (timeToFirstToken.HasValue && timeToFirstToken.Value > 0)
{
activity.SetTag("braintrust.metrics.time_to_first_token", timeToFirstToken.Value);
}
}
}
}
Expand Down
20 changes: 20 additions & 0 deletions tests/Braintrust.Sdk.Tests/Instrumentation/OpenAI/BraintrustOpenAITest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,26 @@ public async Task ChatCompletion_CapturesRequestAndResponse()
var outputNode = JsonNode.Parse(outputJson);
Assert.NotNull(outputNode);
Assert.Contains("The capital of France is Paris", outputJson);

// Verify token metrics were captured
var promptTokens = span.GetTagItem("braintrust.metrics.prompt_tokens");
var completionTokens = span.GetTagItem("braintrust.metrics.completion_tokens");
var totalTokens = span.GetTagItem("braintrust.metrics.tokens");
var timeToFirstToken = span.GetTagItem("braintrust.metrics.time_to_first_token");

Assert.NotNull(promptTokens);
Assert.NotNull(completionTokens);
Assert.NotNull(totalTokens);
Assert.NotNull(timeToFirstToken);

// Verify token counts match the mock response
Assert.Equal(20, Convert.ToInt32(promptTokens));
Assert.Equal(10, Convert.ToInt32(completionTokens));
Assert.Equal(30, Convert.ToInt32(totalTokens));

// Verify time_to_first_token is a non-negative number
var ttft = Convert.ToDouble(timeToFirstToken);
Assert.True(ttft >= 0, "time_to_first_token should be greater than or equal to 0");
}

[Fact]
Expand Down