Skip to content

Commit 472baa2

Browse files
realarkrealark
committed
OAI completions instrumentation to match python and typescript SDKs
1 parent 2ed31d4 commit 472baa2

7 files changed

Lines changed: 306 additions & 83 deletions

File tree

‎src/main/java/dev/braintrust/instrumentation/openai/otel/BraintrustOAISpanAttributes.java‎

Lines changed: 20 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,7 @@
77

88
import com.fasterxml.jackson.databind.ObjectMapper;
99
import com.openai.models.chat.completions.ChatCompletion;
10-
import com.openai.models.chat.completions.ChatCompletionMessage;
1110
import io.opentelemetry.api.trace.Span;
12-
import java.util.List;
1311
import lombok.SneakyThrows;
1412
import lombok.extern.slf4j.Slf4j;
1513

@@ -25,54 +23,34 @@ final class BraintrustOAISpanAttributes {
2523

2624
private BraintrustOAISpanAttributes() {}
2725

28-
/**
29-
* Sets the gen_ai.input.messages attribute with the serialized input messages. This captures
30-
* the user's prompt and system messages before sending to OpenAI.
31-
*/
3226
@SneakyThrows
33-
public static void setInputMessages(Span span, List<?> messages) {
34-
String semconvJson =
35-
GenAiSemconvSerializer.serializeInputMessages(
36-
(List<com.openai.models.chat.completions.ChatCompletionMessageParam>)
37-
messages);
27+
static void setRequestAttributes(
28+
Span span, com.openai.models.chat.completions.ChatCompletionCreateParams request) {
29+
// Set input messages
30+
String semconvJson = GenAiSemconvSerializer.serializeInputMessages(request.messages());
3831
span.setAttribute("gen_ai.input.messages", semconvJson);
39-
}
4032

41-
/**
42-
* Sets the gen_ai.output.messages attribute with the serialized output message. This captures
43-
* the assistant's response from OpenAI for a single choice.
44-
*/
45-
@SneakyThrows
46-
public static void setOutputMessages(
47-
Span span, ChatCompletionMessage message, String finishReason) {
48-
String outputJson = GenAiSemconvSerializer.serializeOutputMessage(message, finishReason);
49-
span.setAttribute("gen_ai.output.messages", outputJson);
50-
}
33+
// Set Braintrust metadata
34+
span.setAttribute("braintrust.metadata.provider", SYSTEM_OPENAI);
5135

52-
/**
53-
* Sets the gen_ai.output.messages attribute for the primary choice in a completion. Logs a
54-
* debug message if there are no choices or multiple choices.
55-
*/
56-
public static void setOutputMessagesFromCompletion(Span span, ChatCompletion completion) {
57-
if (completion.choices().isEmpty()) {
58-
log.debug("no choices in OAI response");
59-
} else if (completion.choices().size() > 1) {
60-
log.debug("multiple choices in OAI response: {}", completion.choices().size());
61-
} else {
62-
// Set gen_ai.output.messages attribute for single choice (most common case)
63-
ChatCompletion.Choice choice = completion.choices().get(0);
64-
setOutputMessages(span, choice.message(), choice.finishReason().toString());
36+
// Set model in metadata if present
37+
try {
38+
var model = request.model();
39+
span.setAttribute("braintrust.metadata.model", model.toString());
40+
} catch (Exception e) {
41+
// If model() throws or returns null, just skip setting it
42+
log.debug("Could not get model from request", e);
6543
}
6644
}
6745

68-
/**
69-
* Sets the braintrust.output_json attribute with a single message. This is used for streaming
70-
* responses to capture output in Braintrust format.
71-
*/
7246
@SneakyThrows
73-
public static void setBraintrustOutputJson(Span span, ChatCompletionMessage message) {
47+
static void setOutputMessagesFromCompletion(Span span, ChatCompletion completion) {
7448
span.setAttribute(
75-
"braintrust.output_json",
76-
JSON_MAPPER.writeValueAsString(new ChatCompletionMessage[] {message}));
49+
"gen_ai.output.messages",
50+
GenAiSemconvSerializer.serializeOutputMessages(completion.choices()));
51+
}
52+
53+
static void setTimeToFirstToken(Span span, double timeInSeconds) {
54+
span.setAttribute("braintrust.metrics.time_to_first_token", timeInSeconds);
7755
}
7856
}

‎src/main/java/dev/braintrust/instrumentation/openai/otel/GenAiSemconvSerializer.java‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,7 @@
77
import com.fasterxml.jackson.databind.ObjectMapper;
88
import com.fasterxml.jackson.databind.SerializerProvider;
99
import com.fasterxml.jackson.databind.module.SimpleModule;
10-
import com.openai.models.chat.completions.ChatCompletionAssistantMessageParam;
11-
import com.openai.models.chat.completions.ChatCompletionContentPartImage;
12-
import com.openai.models.chat.completions.ChatCompletionContentPartText;
13-
import com.openai.models.chat.completions.ChatCompletionDeveloperMessageParam;
14-
import com.openai.models.chat.completions.ChatCompletionMessage;
15-
import com.openai.models.chat.completions.ChatCompletionMessageParam;
16-
import com.openai.models.chat.completions.ChatCompletionMessageToolCall;
17-
import com.openai.models.chat.completions.ChatCompletionSystemMessageParam;
18-
import com.openai.models.chat.completions.ChatCompletionToolMessageParam;
19-
import com.openai.models.chat.completions.ChatCompletionUserMessageParam;
10+
import com.openai.models.chat.completions.*;
2011
import dev.braintrust.trace.Base64Attachment;
2112
import java.io.IOException;
2213
import java.lang.invoke.MethodHandle;
@@ -441,9 +432,12 @@ static String serializeInputMessages(List<ChatCompletionMessageParam> messages)
441432
}
442433

443434
@SneakyThrows
444-
static String serializeOutputMessage(ChatCompletionMessage message, String finishReason) {
445-
SemconvOutputChatMessage outputMessage = transformOutputMessage(message, finishReason);
446-
return JSON_MAPPER.writeValueAsString(new SemconvOutputChatMessage[] {outputMessage});
435+
static String serializeOutputMessages(List<ChatCompletion.Choice> choices) {
436+
var semConvMessages =
437+
choices.stream()
438+
.map(c -> transformOutputMessage(c.message(), c.finishReason().toString()))
439+
.toList();
440+
return JSON_MAPPER.writeValueAsString(semConvMessages);
447441
}
448442

449443
@Nullable

‎src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionService.java‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -100,9 +100,15 @@ private ChatCompletion createWithAttributes(
100100
Context context,
101101
ChatCompletionCreateParams chatCompletionCreateParams,
102102
RequestOptions requestOptions) {
103-
BraintrustOAISpanAttributes.setInputMessages(
104-
Span.current(), chatCompletionCreateParams.messages());
103+
BraintrustOAISpanAttributes.setRequestAttributes(
104+
Span.current(), chatCompletionCreateParams);
105+
106+
long startTimeNanos = System.nanoTime();
105107
ChatCompletion result = delegate.create(chatCompletionCreateParams, requestOptions);
108+
long elapsedNanos = System.nanoTime() - startTimeNanos;
109+
double timeToFirstTokenSeconds = elapsedNanos / 1_000_000_000.0;
110+
111+
BraintrustOAISpanAttributes.setTimeToFirstToken(Span.current(), timeToFirstTokenSeconds);
106112
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(Span.current(), result);
107113
return result;
108114
}
@@ -130,8 +136,9 @@ private StreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
130136
ChatCompletionCreateParams chatCompletionCreateParams,
131137
RequestOptions requestOptions,
132138
boolean newSpan) {
133-
BraintrustOAISpanAttributes.setInputMessages(
134-
Span.current(), chatCompletionCreateParams.messages());
139+
BraintrustOAISpanAttributes.setRequestAttributes(
140+
Span.current(), chatCompletionCreateParams);
141+
long startTimeNanos = System.nanoTime();
135142
StreamResponse<ChatCompletionChunk> result =
136143
delegate.createStreaming(chatCompletionCreateParams, requestOptions);
137144
return new TracingStreamedResponse(
@@ -141,6 +148,7 @@ private StreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
141148
chatCompletionCreateParams,
142149
instrumenter,
143150
captureMessageContent,
144-
newSpan));
151+
newSpan,
152+
startTimeNanos));
145153
}
146154
}

‎src/main/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionServiceAsync.java‎

Lines changed: 16 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -103,14 +103,20 @@ private CompletableFuture<ChatCompletion> createWithAttributes(
103103
Context context,
104104
ChatCompletionCreateParams chatCompletionCreateParams,
105105
RequestOptions requestOptions) {
106-
BraintrustOAISpanAttributes.setInputMessages(
107-
Span.current(), chatCompletionCreateParams.messages());
106+
BraintrustOAISpanAttributes.setRequestAttributes(
107+
Span.current(), chatCompletionCreateParams);
108+
109+
long startTimeNanos = System.nanoTime();
108110
CompletableFuture<ChatCompletion> future =
109111
delegate.create(chatCompletionCreateParams, requestOptions);
110112
future.thenAccept(
111-
r ->
112-
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(
113-
Span.current(), r));
113+
r -> {
114+
long elapsedNanos = System.nanoTime() - startTimeNanos;
115+
double timeToFirstTokenSeconds = elapsedNanos / 1_000_000_000.0;
116+
BraintrustOAISpanAttributes.setTimeToFirstToken(
117+
Span.current(), timeToFirstTokenSeconds);
118+
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(Span.current(), r);
119+
});
114120
return future;
115121
}
116122

@@ -137,8 +143,9 @@ private AsyncStreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
137143
ChatCompletionCreateParams chatCompletionCreateParams,
138144
RequestOptions requestOptions,
139145
boolean newSpan) {
140-
BraintrustOAISpanAttributes.setInputMessages(
141-
Span.current(), chatCompletionCreateParams.messages());
146+
BraintrustOAISpanAttributes.setRequestAttributes(
147+
Span.current(), chatCompletionCreateParams);
148+
long startTimeNanos = System.nanoTime();
142149
AsyncStreamResponse<ChatCompletionChunk> result =
143150
delegate.createStreaming(chatCompletionCreateParams, requestOptions);
144151
return new TracingAsyncStreamedResponse(
@@ -148,6 +155,7 @@ private AsyncStreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
148155
chatCompletionCreateParams,
149156
instrumenter,
150157
captureMessageContent,
151-
newSpan));
158+
newSpan,
159+
startTimeNanos));
152160
}
153161
}

‎src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetryBuilder.java‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@
1313
import io.opentelemetry.api.OpenTelemetry;
1414
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesExtractor;
1515
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiClientMetrics;
16-
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiSpanNameExtractor;
1716
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
1817
import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
18+
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;
1919

2020
/** A builder of {@link OpenAITelemetry}. */
2121
@SuppressWarnings("IdentifierName") // Want to match library's convention
@@ -47,21 +47,23 @@ public OpenAITelemetryBuilder setCaptureMessageContent(boolean captureMessageCon
4747
* OpenAITelemetryBuilder}.
4848
*/
4949
public OpenAITelemetry build() {
50+
// Use hardcoded span names to match Python/TypeScript SDKs
51+
SpanNameExtractor<ChatCompletionCreateParams> chatSpanNameExtractor =
52+
request -> "Chat Completion";
53+
SpanNameExtractor<EmbeddingCreateParams> embeddingSpanNameExtractor =
54+
request -> "Embedding";
55+
5056
Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter =
5157
Instrumenter.<ChatCompletionCreateParams, ChatCompletion>builder(
52-
openTelemetry,
53-
INSTRUMENTATION_NAME,
54-
GenAiSpanNameExtractor.create(ChatAttributesGetter.INSTANCE))
58+
openTelemetry, INSTRUMENTATION_NAME, chatSpanNameExtractor)
5559
.addAttributesExtractor(
5660
GenAiAttributesExtractor.create(ChatAttributesGetter.INSTANCE))
5761
.addOperationMetrics(GenAiClientMetrics.get())
5862
.buildInstrumenter();
5963

6064
Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingsInstrumenter =
6165
Instrumenter.<EmbeddingCreateParams, CreateEmbeddingResponse>builder(
62-
openTelemetry,
63-
INSTRUMENTATION_NAME,
64-
GenAiSpanNameExtractor.create(EmbeddingAttributesGetter.INSTANCE))
66+
openTelemetry, INSTRUMENTATION_NAME, embeddingSpanNameExtractor)
6567
.addAttributesExtractor(
6668
GenAiAttributesExtractor.create(EmbeddingAttributesGetter.INSTANCE))
6769
.addOperationMetrics(GenAiClientMetrics.get())

‎src/main/java/dev/braintrust/instrumentation/openai/otel/StreamListener.java‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ final class StreamListener {
2828
private final boolean captureMessageContent;
2929
private final boolean newSpan;
3030
private final AtomicBoolean hasEnded;
31+
private final AtomicBoolean firstChunkReceived;
32+
private final long startTimeNanos;
3133

3234
@Nullable private CompletionUsage usage;
3335
@Nullable private String model;
@@ -38,18 +40,29 @@ final class StreamListener {
3840
ChatCompletionCreateParams request,
3941
Instrumenter<ChatCompletionCreateParams, ChatCompletion> instrumenter,
4042
boolean captureMessageContent,
41-
boolean newSpan) {
43+
boolean newSpan,
44+
long startTimeNanos) {
4245
this.context = context;
4346
this.request = request;
4447
this.instrumenter = instrumenter;
4548
this.captureMessageContent = captureMessageContent;
4649
this.newSpan = newSpan;
50+
this.startTimeNanos = startTimeNanos;
4751
choiceBuffers = new ArrayList<>();
4852
hasEnded = new AtomicBoolean();
53+
firstChunkReceived = new AtomicBoolean();
4954
}
5055

5156
@SneakyThrows
5257
void onChunk(ChatCompletionChunk chunk) {
58+
// Calculate time to first token on the first chunk
59+
if (firstChunkReceived.compareAndSet(false, true)) {
60+
long elapsedNanos = System.nanoTime() - startTimeNanos;
61+
double timeToFirstTokenSeconds = elapsedNanos / 1_000_000_000.0;
62+
BraintrustOAISpanAttributes.setTimeToFirstToken(
63+
Span.fromContext(context), timeToFirstTokenSeconds);
64+
}
65+
5366
model = chunk.model();
5467
responseId = chunk.id();
5568
chunk.usage().ifPresent(u -> usage = u);
@@ -66,8 +79,6 @@ void onChunk(ChatCompletionChunk chunk) {
6679
buffer.append(choice.delta());
6780
if (choice.finishReason().isPresent()) {
6881
buffer.finishReason = choice.finishReason().get().toString();
69-
BraintrustOAISpanAttributes.setBraintrustOutputJson(
70-
Span.fromContext(context), buffer.toChoice().message());
7182
}
7283
}
7384
}
@@ -102,7 +113,10 @@ void endSpan(@Nullable Throwable error) {
102113
}
103114

104115
if (newSpan) {
105-
instrumenter.end(context, request, result.build(), error);
116+
ChatCompletion completion = result.build();
117+
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(
118+
Span.fromContext(context), completion);
119+
instrumenter.end(context, request, completion, error);
106120
}
107121
}
108122
}

0 commit comments

Comments
 (0)