Skip to content

Commit 96e287e

Browse files
realarkrealark
committed
wip -- tests passing
1 parent 65eae57 commit 96e287e

4 files changed

Lines changed: 52 additions & 74 deletions

File tree

‎examples/src/main/java/dev/braintrust/examples/OpenAIInstrumentationExample.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,9 @@ private static void chatCompletionsExample(OpenAIClient openAIClient) {
5555
// Create text content part
5656
ChatCompletionContentPartText textPart =
5757
ChatCompletionContentPartText.builder()
58-
.text("What city and landmark are shown in this travel poster? Describe the image.")
58+
.text(
59+
"What city and landmark are shown in this travel poster? Describe"
60+
+ " the image.")
5961
.build();
6062
ChatCompletionContentPart textContentPart = ChatCompletionContentPart.ofText(textPart);
6163

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

Lines changed: 25 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,13 @@
77

88
import static io.opentelemetry.api.common.AttributeKey.stringKey;
99

10+
import com.fasterxml.jackson.annotation.JsonInclude;
11+
import com.fasterxml.jackson.annotation.JsonProperty;
1012
import com.fasterxml.jackson.core.JsonGenerator;
1113
import com.fasterxml.jackson.databind.JsonSerializer;
1214
import com.fasterxml.jackson.databind.ObjectMapper;
1315
import com.fasterxml.jackson.databind.SerializerProvider;
1416
import com.fasterxml.jackson.databind.module.SimpleModule;
15-
import com.fasterxml.jackson.annotation.JsonInclude;
16-
import com.fasterxml.jackson.annotation.JsonProperty;
1717
import com.openai.models.chat.completions.ChatCompletion;
1818
import com.openai.models.chat.completions.ChatCompletionAssistantMessageParam;
1919
import com.openai.models.chat.completions.ChatCompletionContentPartImage;
@@ -419,26 +419,28 @@ private static SemconvOutputChatMessage transformOutputMessage(
419419

420420
// Handle text content
421421
message.content()
422-
.ifPresent(content -> {
423-
if (!content.isEmpty()) {
424-
parts.add(new TextPart(content));
425-
}
426-
});
422+
.ifPresent(
423+
content -> {
424+
if (!content.isEmpty()) {
425+
parts.add(new TextPart(content));
426+
}
427+
});
427428

428429
// Handle tool calls
429430
message.toolCalls()
430-
.ifPresent(toolCalls -> {
431-
for (var toolCall : toolCalls) {
432-
FunctionAccess functionAccess = getFunctionAccess(toolCall);
433-
if (functionAccess != null) {
434-
parts.add(
435-
new ToolCallRequestPart(
436-
functionAccess.name(),
437-
functionAccess.id(),
438-
functionAccess.arguments()));
439-
}
440-
}
441-
});
431+
.ifPresent(
432+
toolCalls -> {
433+
for (var toolCall : toolCalls) {
434+
FunctionAccess functionAccess = getFunctionAccess(toolCall);
435+
if (functionAccess != null) {
436+
parts.add(
437+
new ToolCallRequestPart(
438+
functionAccess.name(),
439+
functionAccess.id(),
440+
functionAccess.arguments()));
441+
}
442+
}
443+
});
442444

443445
// The role from ChatCompletionMessage is always "assistant" for output messages
444446
return new SemconvOutputChatMessage("assistant", parts, finishReason);
@@ -540,10 +542,10 @@ public static void emitCompletionLogEvents(
540542
} else {
541543
// Set gen_ai.output.messages attribute for single choice (most common case)
542544
ChatCompletion.Choice choice = completion.choices().get(0);
543-
SemconvOutputChatMessage outputMessage = transformOutputMessage(
544-
choice.message(),
545-
choice.finishReason().toString());
546-
String outputJson = JSON_MAPPER.writeValueAsString(new SemconvOutputChatMessage[] {outputMessage});
545+
SemconvOutputChatMessage outputMessage =
546+
transformOutputMessage(choice.message(), choice.finishReason().toString());
547+
String outputJson =
548+
JSON_MAPPER.writeValueAsString(new SemconvOutputChatMessage[] {outputMessage});
547549
Span.current().setAttribute("gen_ai.output.messages", outputJson);
548550
}
549551
for (ChatCompletion.Choice choice : completion.choices()) {

‎src/main/java/dev/braintrust/trace/Base64Attachment.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,12 @@
2525
*/
2626
public class Base64Attachment {
2727
@JsonProperty("type")
28-
@Getter private final String type = "base64_attachment";
28+
@Getter
29+
private final String type = "base64_attachment";
30+
2931
@JsonProperty("content")
30-
@Getter private final String base64Data;
32+
@Getter
33+
private final String base64Data;
3134

3235
private Base64Attachment(@Nonnull String base64Data) {
3336
if (Objects.requireNonNull(base64Data).isEmpty()) {

‎src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java‎

Lines changed: 19 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -146,10 +146,11 @@ void testWrapOpenAi() {
146146
"chatcmpl-test123",
147147
span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.id")));
148148
assertEquals(
149-
"[{\"content\":\"You are a helpful"
150-
+ " assistant\",\"role\":\"system\",\"valid\":true},{\"content\":\"What is the"
151-
+ " capital of France?\",\"role\":\"user\",\"valid\":true}]",
152-
span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")));
149+
"[{\"role\":\"system\",\"parts\":[{\"type\":\"text\",\"content\":\"You are a"
150+
+ " helpful"
151+
+ " assistant\"}]},{\"role\":\"user\",\"parts\":[{\"type\":\"text\",\"content\":\"What"
152+
+ " is the capital of France?\"}]}]",
153+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.input.messages")));
153154
assertEquals(
154155
"project_name:unit-test-project",
155156
span.getAttributes().get(AttributeKey.stringKey("braintrust.parent")));
@@ -162,12 +163,11 @@ void testWrapOpenAi() {
162163
span.getAttributes().get(AttributeKey.doubleKey("gen_ai.request.temperature")));
163164

164165
String outputJson =
165-
span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json"));
166-
assertNotNull(outputJson);
167-
var outputMessages = JSON_MAPPER.readTree(outputJson);
168-
assertEquals(1, outputMessages.size());
169-
var messageZero = outputMessages.get(0);
170-
assertEquals("The capital of France is Paris.", messageZero.get("content").asText());
166+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.output.messages"));
167+
assertEquals(
168+
"[{\"role\":\"assistant\",\"parts\":[{\"type\":\"text\",\"content\":\"The capital"
169+
+ " of France is Paris.\"}],\"finish_reason\":\"stop\"}]",
170+
outputJson);
171171

172172
assertEquals(
173173
"chatcmpl-test123",
@@ -432,41 +432,15 @@ void testWrapOpenAiWithImageAttachment() {
432432

433433
// Verify input JSON captures both text and image content
434434
String inputJson =
435-
span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json"));
435+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.input.messages"));
436+
assertEquals(
437+
"[{\"role\":\"system\",\"parts\":[{\"type\":\"text\",\"content\":\"You are a helpful assistant that can analyze images\"}]},{\"role\":\"user\",\"parts\":[{\"type\":\"text\",\"content\":\"What's in this image?\"},{\"type\":\"base64_attachment\",\"content\":\"%s\"}]}]"
438+
.formatted(imageDataUrl),
439+
inputJson);
436440
assertNotNull(inputJson);
437441
var inputMessages = JSON_MAPPER.readTree(inputJson);
438442
assertEquals(2, inputMessages.size()); // system message + user message
439443

440-
// Verify system message
441-
var systemMessage = inputMessages.get(0);
442-
assertEquals("system", systemMessage.get("role").asText());
443-
assertEquals(
444-
"You are a helpful assistant that can analyze images",
445-
systemMessage.get("content").asText());
446-
447-
// Verify user message with image
448-
var userMsg = inputMessages.get(1);
449-
assertEquals("user", userMsg.get("role").asText());
450-
assertTrue(userMsg.has("content"));
451-
var content = userMsg.get("content");
452-
assertTrue(content.isArray());
453-
assertEquals(2, content.size()); // text + image
454-
455-
// Verify text content part
456-
var textContent = content.get(0);
457-
assertEquals("text", textContent.get("type").asText());
458-
assertEquals("What's in this image?", textContent.get("text").asText());
459-
460-
// Verify image content part (now serialized as base64_attachment)
461-
var imageContent = content.get(1);
462-
assertEquals("base64_attachment", imageContent.get("type").asText());
463-
assertTrue(imageContent.has("content"));
464-
assertTrue(imageContent.get("content").asText().startsWith("data:image/jpeg;base64,"));
465-
assertEquals(
466-
imageDataUrl,
467-
imageContent.get("content").asText(),
468-
"base64 data not correctly serialized");
469-
470444
// Verify usage metrics
471445
assertEquals(
472446
150L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.input_tokens")));
@@ -475,13 +449,10 @@ void testWrapOpenAiWithImageAttachment() {
475449

476450
// Verify output JSON
477451
String outputJson =
478-
span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json"));
479-
assertNotNull(outputJson);
480-
var outputMessages = JSON_MAPPER.readTree(outputJson);
481-
assertEquals(1, outputMessages.size());
482-
var messageZero = outputMessages.get(0);
452+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.output.messages"));
483453
assertEquals(
484-
"This image shows the Eiffel Tower in Paris, France.",
485-
messageZero.get("content").asText());
454+
"[{\"role\":\"assistant\",\"parts\":[{\"type\":\"text\",\"content\":\"This image"
455+
+ " shows the Eiffel Tower in Paris, France.\"}],\"finish_reason\":\"stop\"}]",
456+
outputJson);
486457
}
487458
}

0 commit comments

Comments
 (0)