diff --git a/src/Braintrust.Sdk/Instrumentation/OpenAI/InstrumentedOpenAIClient.cs b/src/Braintrust.Sdk/Instrumentation/OpenAI/InstrumentedOpenAIClient.cs index 234a95d..23edc1b 100644 --- a/src/Braintrust.Sdk/Instrumentation/OpenAI/InstrumentedOpenAIClient.cs +++ b/src/Braintrust.Sdk/Instrumentation/OpenAI/InstrumentedOpenAIClient.cs @@ -147,16 +147,20 @@ public override ClientResult CompleteChat(IEnumerable> 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; @@ -219,7 +227,7 @@ public override async Task> 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"); { @@ -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(); + var completionTokens = usage["completion_tokens"]?.GetValue(); + var totalTokens = usage["total_tokens"]?.GetValue(); + + 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); + } } } } diff --git a/tests/Braintrust.Sdk.Tests/Instrumentation/OpenAI/BraintrustOpenAITest.cs b/tests/Braintrust.Sdk.Tests/Instrumentation/OpenAI/BraintrustOpenAITest.cs index dabd3e7..febc5d2 100644 --- a/tests/Braintrust.Sdk.Tests/Instrumentation/OpenAI/BraintrustOpenAITest.cs +++ b/tests/Braintrust.Sdk.Tests/Instrumentation/OpenAI/BraintrustOpenAITest.cs @@ -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]