Skip to content

enhance OAI completions instrumentation #10

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
Andrew Kent (realark) merged 1 commit into from
Nov 6, 2025
Merged

enhance OAI completions instrumentation #10

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
62 changes: 20 additions & 42 deletions src/main/java/dev/braintrust/instrumentation/openai/otel/BraintrustOAISpanAttributes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,7 @@

import com.fasterxml.jackson.databind.ObjectMapper;
import com.openai.models.chat.completions.ChatCompletion;
import com.openai.models.chat.completions.ChatCompletionMessage;
import io.opentelemetry.api.trace.Span;
import java.util.List;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

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

private BraintrustOAISpanAttributes() {}

/**
* Sets the gen_ai.input.messages attribute with the serialized input messages. This captures
* the user's prompt and system messages before sending to OpenAI.
*/
@SneakyThrows
public static void setInputMessages(Span span, List<?> messages) {
String semconvJson =
GenAiSemconvSerializer.serializeInputMessages(
(List<com.openai.models.chat.completions.ChatCompletionMessageParam>)
messages);
static void setRequestAttributes(
Span span, com.openai.models.chat.completions.ChatCompletionCreateParams request) {
// Set input messages
String semconvJson = GenAiSemconvSerializer.serializeInputMessages(request.messages());
span.setAttribute("gen_ai.input.messages", semconvJson);
}

/**
* Sets the gen_ai.output.messages attribute with the serialized output message. This captures
* the assistant's response from OpenAI for a single choice.
*/
@SneakyThrows
public static void setOutputMessages(
Span span, ChatCompletionMessage message, String finishReason) {
String outputJson = GenAiSemconvSerializer.serializeOutputMessage(message, finishReason);
span.setAttribute("gen_ai.output.messages", outputJson);
}
// Set Braintrust metadata
span.setAttribute("braintrust.metadata.provider", SYSTEM_OPENAI);

/**
* Sets the gen_ai.output.messages attribute for the primary choice in a completion. Logs a
* debug message if there are no choices or multiple choices.
*/
public static void setOutputMessagesFromCompletion(Span span, ChatCompletion completion) {
if (completion.choices().isEmpty()) {
log.debug("no choices in OAI response");
} else if (completion.choices().size() > 1) {
log.debug("multiple choices in OAI response: {}", completion.choices().size());
} else {
// Set gen_ai.output.messages attribute for single choice (most common case)
ChatCompletion.Choice choice = completion.choices().get(0);
setOutputMessages(span, choice.message(), choice.finishReason().toString());
// Set model in metadata if present
try {
var model = request.model();
span.setAttribute("braintrust.metadata.model", model.toString());
} catch (Exception e) {
// If model() throws or returns null, just skip setting it
log.debug("Could not get model from request", e);
}
}

/**
* Sets the braintrust.output_json attribute with a single message. This is used for streaming
* responses to capture output in Braintrust format.
*/
@SneakyThrows
public static void setBraintrustOutputJson(Span span, ChatCompletionMessage message) {
static void setOutputMessagesFromCompletion(Span span, ChatCompletion completion) {
span.setAttribute(
"braintrust.output_json",
JSON_MAPPER.writeValueAsString(new ChatCompletionMessage[] {message}));
"gen_ai.output.messages",
GenAiSemconvSerializer.serializeOutputMessages(completion.choices()));
}

static void setTimeToFirstToken(Span span, double timeInSeconds) {
span.setAttribute("braintrust.metrics.time_to_first_token", timeInSeconds);
}
}
20 changes: 7 additions & 13 deletions src/main/java/dev/braintrust/instrumentation/openai/otel/GenAiSemconvSerializer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,7 @@
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializerProvider;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.openai.models.chat.completions.ChatCompletionAssistantMessageParam;
import com.openai.models.chat.completions.ChatCompletionContentPartImage;
import com.openai.models.chat.completions.ChatCompletionContentPartText;
import com.openai.models.chat.completions.ChatCompletionDeveloperMessageParam;
import com.openai.models.chat.completions.ChatCompletionMessage;
import com.openai.models.chat.completions.ChatCompletionMessageParam;
import com.openai.models.chat.completions.ChatCompletionMessageToolCall;
import com.openai.models.chat.completions.ChatCompletionSystemMessageParam;
import com.openai.models.chat.completions.ChatCompletionToolMessageParam;
import com.openai.models.chat.completions.ChatCompletionUserMessageParam;
import com.openai.models.chat.completions.*;
import dev.braintrust.trace.Base64Attachment;
import java.io.IOException;
import java.lang.invoke.MethodHandle;
Expand Down Expand Up @@ -441,9 +432,12 @@ static String serializeInputMessages(List<ChatCompletionMessageParam> messages)
}

@SneakyThrows
static String serializeOutputMessage(ChatCompletionMessage message, String finishReason) {
SemconvOutputChatMessage outputMessage = transformOutputMessage(message, finishReason);
return JSON_MAPPER.writeValueAsString(new SemconvOutputChatMessage[] {outputMessage});
static String serializeOutputMessages(List<ChatCompletion.Choice> choices) {
var semConvMessages =
choices.stream()
.map(c -> transformOutputMessage(c.message(), c.finishReason().toString()))
.toList();
return JSON_MAPPER.writeValueAsString(semConvMessages);
}

@Nullable
Expand Down
18 changes: 13 additions & 5 deletions ...in/java/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,9 +100,15 @@ private ChatCompletion createWithAttributes(
Context context,
ChatCompletionCreateParams chatCompletionCreateParams,
RequestOptions requestOptions) {
BraintrustOAISpanAttributes.setInputMessages(
Span.current(), chatCompletionCreateParams.messages());
BraintrustOAISpanAttributes.setRequestAttributes(
Span.current(), chatCompletionCreateParams);

long startTimeNanos = System.nanoTime();
ChatCompletion result = delegate.create(chatCompletionCreateParams, requestOptions);
long elapsedNanos = System.nanoTime() - startTimeNanos;
double timeToFirstTokenSeconds = elapsedNanos / 1_000_000_000.0;

BraintrustOAISpanAttributes.setTimeToFirstToken(Span.current(), timeToFirstTokenSeconds);
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(Span.current(), result);
return result;
}
Expand Down Expand Up @@ -130,8 +136,9 @@ private StreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
ChatCompletionCreateParams chatCompletionCreateParams,
RequestOptions requestOptions,
boolean newSpan) {
BraintrustOAISpanAttributes.setInputMessages(
Span.current(), chatCompletionCreateParams.messages());
BraintrustOAISpanAttributes.setRequestAttributes(
Span.current(), chatCompletionCreateParams);
long startTimeNanos = System.nanoTime();
StreamResponse<ChatCompletionChunk> result =
delegate.createStreaming(chatCompletionCreateParams, requestOptions);
return new TracingStreamedResponse(
Expand All @@ -141,6 +148,7 @@ private StreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
chatCompletionCreateParams,
instrumenter,
captureMessageContent,
newSpan));
newSpan,
startTimeNanos));
}
}
24 changes: 16 additions & 8 deletions ...va/dev/braintrust/instrumentation/openai/otel/InstrumentedChatCompletionServiceAsync.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,14 +103,20 @@ private CompletableFuture<ChatCompletion> createWithAttributes(
Context context,
ChatCompletionCreateParams chatCompletionCreateParams,
RequestOptions requestOptions) {
BraintrustOAISpanAttributes.setInputMessages(
Span.current(), chatCompletionCreateParams.messages());
BraintrustOAISpanAttributes.setRequestAttributes(
Span.current(), chatCompletionCreateParams);

long startTimeNanos = System.nanoTime();
CompletableFuture<ChatCompletion> future =
delegate.create(chatCompletionCreateParams, requestOptions);
future.thenAccept(
r ->
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(
Span.current(), r));
r -> {
long elapsedNanos = System.nanoTime() - startTimeNanos;
double timeToFirstTokenSeconds = elapsedNanos / 1_000_000_000.0;
BraintrustOAISpanAttributes.setTimeToFirstToken(
Span.current(), timeToFirstTokenSeconds);
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(Span.current(), r);
});
return future;
}

Expand All @@ -137,8 +143,9 @@ private AsyncStreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
ChatCompletionCreateParams chatCompletionCreateParams,
RequestOptions requestOptions,
boolean newSpan) {
BraintrustOAISpanAttributes.setInputMessages(
Span.current(), chatCompletionCreateParams.messages());
BraintrustOAISpanAttributes.setRequestAttributes(
Span.current(), chatCompletionCreateParams);
long startTimeNanos = System.nanoTime();
AsyncStreamResponse<ChatCompletionChunk> result =
delegate.createStreaming(chatCompletionCreateParams, requestOptions);
return new TracingAsyncStreamedResponse(
Expand All @@ -148,6 +155,7 @@ private AsyncStreamResponse<ChatCompletionChunk> createStreamingWithAttributes(
chatCompletionCreateParams,
instrumenter,
captureMessageContent,
newSpan));
newSpan,
startTimeNanos));
}
}
16 changes: 9 additions & 7 deletions src/main/java/dev/braintrust/instrumentation/openai/otel/OpenAITelemetryBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiAttributesExtractor;
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiClientMetrics;
import io.opentelemetry.instrumentation.api.incubator.semconv.genai.GenAiSpanNameExtractor;
import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
import io.opentelemetry.instrumentation.api.instrumenter.SpanNameExtractor;

/** A builder of {@link OpenAITelemetry}. */
@SuppressWarnings("IdentifierName") // Want to match library's convention
Expand Down Expand Up @@ -47,21 +47,23 @@ public OpenAITelemetryBuilder setCaptureMessageContent(boolean captureMessageCon
* OpenAITelemetryBuilder}.
*/
public OpenAITelemetry build() {
// Use hardcoded span names to match Python/TypeScript SDKs
SpanNameExtractor<ChatCompletionCreateParams> chatSpanNameExtractor =
request -> "Chat Completion";
SpanNameExtractor<EmbeddingCreateParams> embeddingSpanNameExtractor =
request -> "Embedding";

Instrumenter<ChatCompletionCreateParams, ChatCompletion> chatInstrumenter =
Instrumenter.<ChatCompletionCreateParams, ChatCompletion>builder(
openTelemetry,
INSTRUMENTATION_NAME,
GenAiSpanNameExtractor.create(ChatAttributesGetter.INSTANCE))
openTelemetry, INSTRUMENTATION_NAME, chatSpanNameExtractor)
.addAttributesExtractor(
GenAiAttributesExtractor.create(ChatAttributesGetter.INSTANCE))
.addOperationMetrics(GenAiClientMetrics.get())
.buildInstrumenter();

Instrumenter<EmbeddingCreateParams, CreateEmbeddingResponse> embeddingsInstrumenter =
Instrumenter.<EmbeddingCreateParams, CreateEmbeddingResponse>builder(
openTelemetry,
INSTRUMENTATION_NAME,
GenAiSpanNameExtractor.create(EmbeddingAttributesGetter.INSTANCE))
openTelemetry, INSTRUMENTATION_NAME, embeddingSpanNameExtractor)
.addAttributesExtractor(
GenAiAttributesExtractor.create(EmbeddingAttributesGetter.INSTANCE))
.addOperationMetrics(GenAiClientMetrics.get())
Expand Down
22 changes: 18 additions & 4 deletions src/main/java/dev/braintrust/instrumentation/openai/otel/StreamListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@ final class StreamListener {
private final boolean captureMessageContent;
private final boolean newSpan;
private final AtomicBoolean hasEnded;
private final AtomicBoolean firstChunkReceived;
private final long startTimeNanos;

@Nullable private CompletionUsage usage;
@Nullable private String model;
Expand All @@ -38,18 +40,29 @@ final class StreamListener {
ChatCompletionCreateParams request,
Instrumenter<ChatCompletionCreateParams, ChatCompletion> instrumenter,
boolean captureMessageContent,
boolean newSpan) {
boolean newSpan,
long startTimeNanos) {
this.context = context;
this.request = request;
this.instrumenter = instrumenter;
this.captureMessageContent = captureMessageContent;
this.newSpan = newSpan;
this.startTimeNanos = startTimeNanos;
choiceBuffers = new ArrayList<>();
hasEnded = new AtomicBoolean();
firstChunkReceived = new AtomicBoolean();
}

@SneakyThrows
void onChunk(ChatCompletionChunk chunk) {
// Calculate time to first token on the first chunk
if (firstChunkReceived.compareAndSet(false, true)) {
long elapsedNanos = System.nanoTime() - startTimeNanos;
double timeToFirstTokenSeconds = elapsedNanos / 1_000_000_000.0;
BraintrustOAISpanAttributes.setTimeToFirstToken(
Span.fromContext(context), timeToFirstTokenSeconds);
}

model = chunk.model();
responseId = chunk.id();
chunk.usage().ifPresent(u -> usage = u);
Expand All @@ -66,8 +79,6 @@ void onChunk(ChatCompletionChunk chunk) {
buffer.append(choice.delta());
if (choice.finishReason().isPresent()) {
buffer.finishReason = choice.finishReason().get().toString();
BraintrustOAISpanAttributes.setBraintrustOutputJson(
Span.fromContext(context), buffer.toChoice().message());
}
}
}
Expand Down Expand Up @@ -102,7 +113,10 @@ void endSpan(@Nullable Throwable error) {
}

if (newSpan) {
instrumenter.end(context, request, result.build(), error);
ChatCompletion completion = result.build();
BraintrustOAISpanAttributes.setOutputMessagesFromCompletion(
Span.fromContext(context), completion);
instrumenter.end(context, request, completion, error);
}
}
}
Loading