diff --git a/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java b/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java index 0814c16..ce25f4d 100644 --- a/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java +++ b/src/main/java/dev/braintrust/instrumentation/openai/otel/ChatCompletionEventsHelper.java @@ -7,9 +7,14 @@ import static io.opentelemetry.api.common.AttributeKey.stringKey; +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.databind.JsonSerializer; 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.ChatCompletion; 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.ChatCompletionCreateParams; import com.openai.models.chat.completions.ChatCompletionDeveloperMessageParam; @@ -18,12 +23,14 @@ import com.openai.models.chat.completions.ChatCompletionSystemMessageParam; import com.openai.models.chat.completions.ChatCompletionToolMessageParam; import com.openai.models.chat.completions.ChatCompletionUserMessageParam; +import dev.braintrust.trace.Base64Attachment; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.common.Value; import io.opentelemetry.api.logs.LogRecordBuilder; import io.opentelemetry.api.logs.Logger; import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Context; +import java.io.IOException; import java.lang.invoke.MethodHandle; import java.lang.invoke.MethodHandles; import java.lang.invoke.MethodType; @@ -41,8 +48,38 @@ final class ChatCompletionEventsHelper { private static final AttributeKey EVENT_NAME = stringKey("event.name"); - private static final ObjectMapper JSON_MAPPER = - new com.fasterxml.jackson.databind.ObjectMapper(); + private static final ObjectMapper JSON_MAPPER = createObjectMapper(); + + private static ObjectMapper createObjectMapper() { + final JsonSerializer attachmentSerializer = + Base64Attachment.createSerializer(); + ObjectMapper mapper = new ObjectMapper(); + SimpleModule module = new SimpleModule(); + module.addSerializer( + ChatCompletionContentPartImage.class, + new JsonSerializer<>() { + @Override + public void serialize( + ChatCompletionContentPartImage value, + JsonGenerator gen, + SerializerProvider serializers) + throws IOException { + try { + var attachment = + Base64Attachment.of( + value.validate().imageUrl().validate().url()); + attachmentSerializer.serialize(attachment, gen, serializers); + } catch (Exception e) { + JsonSerializer defaultSerializer = + serializers.findValueSerializer( + ChatCompletionContentPartImage.class, null); + defaultSerializer.serialize(value, gen, serializers); + } + } + }); + mapper.registerModule(module); + return mapper; + } @SneakyThrows public static void emitPromptLogEvents( diff --git a/src/main/java/dev/braintrust/trace/Base64Attachment.java b/src/main/java/dev/braintrust/trace/Base64Attachment.java new file mode 100644 index 0000000..84d4f7c --- /dev/null +++ b/src/main/java/dev/braintrust/trace/Base64Attachment.java @@ -0,0 +1,121 @@ +package dev.braintrust.trace; + +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.databind.JsonSerializer; +import com.fasterxml.jackson.databind.SerializerProvider; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Base64; +import java.util.Objects; +import javax.annotation.Nonnull; +import lombok.Getter; + +/** + * Utility to serialize LLM attachment data in a braintrust-friendly manner. + * + *

Users of the SDK likely don't need to use this utility directly because instrumentation will + * properly serialize messages out of the box. + * + *

The serialized json will conform to the otel input/output GenericPart schema. See + * https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-input-messages.json and + * https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-output-messages.json + */ +public class Base64Attachment { + @Getter private final String type = "base64_attachment"; + @Getter private final String base64Data; + + private Base64Attachment(@Nonnull String base64Data) { + if (Objects.requireNonNull(base64Data).isEmpty()) { + throw new IllegalArgumentException("base64Data cannot be empty"); + } + // Check for data URL prefix (e.g., "data:image/png;base64,...") + if (!base64Data.startsWith("data:") || !base64Data.contains(";base64,")) { + throw new IllegalArgumentException( + "base64Data must be a data URL with format:" + + " data:;base64,"); + } + + this.base64Data = base64Data; + } + + /** + * Create a new attachment out of base64 data + * + * @param base64DataUri must conform to data:(content-type);base64,BYTES + */ + public static Base64Attachment of(String base64DataUri) { + return new Base64Attachment(base64DataUri); + } + + /** convenience utility to convert a file to a base64 attachment */ + public static Base64Attachment ofFile(ContentType contentType, String filePath) { + try { + Path path = Paths.get(filePath); + byte[] fileBytes = Files.readAllBytes(path); + String base64Encoded = Base64.getEncoder().encodeToString(fileBytes); + String dataUrl = "data:" + contentType.getMimeType() + ";base64," + base64Encoded; + return of(dataUrl); + } catch (IOException e) { + throw new RuntimeException("Failed to read file: " + filePath, e); + } + } + + /** create a jackson serializer for attachment data */ + public static JsonSerializer createSerializer() { + return new JsonSerializer<>() { + @Override + public void serialize( + Base64Attachment value, JsonGenerator gen, SerializerProvider serializers) + throws IOException { + gen.writeStartObject(); + try { + gen.writeStringField("type", value.type); + gen.writeStringField("content", value.base64Data); + } finally { + gen.writeEndObject(); + } + } + }; + } + + public static class ContentType { + // Common image formats + public static ContentType IMAGE_PNG = new ContentType("image/png"); + public static ContentType IMAGE_JPEG = new ContentType("image/jpeg"); + public static ContentType IMAGE_GIF = new ContentType("image/gif"); + public static ContentType IMAGE_WEBP = new ContentType("image/webp"); + public static ContentType IMAGE_SVG = new ContentType("image/svg+xml"); + + // Common document formats + public static ContentType APPLICATION_PDF = new ContentType("application/pdf"); + public static ContentType TEXT_PLAIN = new ContentType("text/plain"); + public static ContentType APPLICATION_JSON = new ContentType("application/json"); + + public static ContentType of(@Nonnull String mimeType) { + return new ContentType(mimeType); + } + + @Getter private final @Nonnull String mimeType; + + @Override + public int hashCode() { + return mimeType.hashCode(); + } + + @Override + public boolean equals(Object obj) { + if (obj instanceof ContentType) { + return mimeType.equals(((ContentType) obj).mimeType); + } else { + return super.equals(obj); + } + } + + private ContentType(@Nonnull String mimeType) { + Objects.requireNonNull(mimeType); + this.mimeType = mimeType.toLowerCase(); + } + } +} diff --git a/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java b/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java index 434adab..08bb42e 100644 --- a/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java +++ b/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java @@ -10,13 +10,19 @@ import com.openai.client.OpenAIClient; import com.openai.client.okhttp.OpenAIOkHttpClient; import com.openai.models.ChatModel; +import com.openai.models.chat.completions.ChatCompletionContentPart; +import com.openai.models.chat.completions.ChatCompletionContentPartImage; +import com.openai.models.chat.completions.ChatCompletionContentPartText; import com.openai.models.chat.completions.ChatCompletionCreateParams; import com.openai.models.chat.completions.ChatCompletionStreamOptions; +import com.openai.models.chat.completions.ChatCompletionUserMessageParam; import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.trace.Base64Attachment; import dev.braintrust.trace.BraintrustTracing; import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.sdk.OpenTelemetrySdk; +import java.util.Arrays; import java.util.concurrent.TimeUnit; import lombok.SneakyThrows; import org.junit.jupiter.api.BeforeEach; @@ -295,4 +301,187 @@ void testWrapOpenAiStreaming() { var messageZero = outputMessages.get(0); assertEquals("The capital of France is Paris.", messageZero.get("content").asText()); } + + @Test + @SneakyThrows + void testWrapOpenAiWithImageAttachment() { + // Mock the OpenAI API response for vision request + wireMock.stubFor( + post(urlEqualTo("/chat/completions")) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "id": "chatcmpl-test456", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "This image shows the Eiffel Tower in Paris, France." + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 150, + "completion_tokens": 15, + "total_tokens": 165 + } + } + """))); + + var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true); + + // Create OpenAI client pointing to WireMock server + OpenAIClient openAIClient = + OpenAIOkHttpClient.builder() + .baseUrl("http://localhost:" + wireMock.getPort()) + .apiKey("test-api-key") + .build(); + + // Wrap with Braintrust instrumentation + openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, openAIClient); + + String imageDataUrl = + Base64Attachment.ofFile( + Base64Attachment.ContentType.IMAGE_JPEG, + "src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg") + .getBase64Data(); + + // Create text content part + ChatCompletionContentPartText textPart = + ChatCompletionContentPartText.builder().text("What's in this image?").build(); + ChatCompletionContentPart textContentPart = ChatCompletionContentPart.ofText(textPart); + + // Create image content part with base64-encoded image + ChatCompletionContentPartImage imagePart = + ChatCompletionContentPartImage.builder() + .imageUrl( + ChatCompletionContentPartImage.ImageUrl.builder() + // .url("https://example.com/eiffel-tower.jpg") + .url(imageDataUrl) + .detail(ChatCompletionContentPartImage.ImageUrl.Detail.HIGH) + .build()) + .build(); + ChatCompletionContentPart imageContentPart = + ChatCompletionContentPart.ofImageUrl(imagePart); + + // Create user message with both text and image + ChatCompletionUserMessageParam userMessage = + ChatCompletionUserMessageParam.builder() + .contentOfArrayOfContentParts( + Arrays.asList(textContentPart, imageContentPart)) + .build(); + + var request = + ChatCompletionCreateParams.builder() + .model(ChatModel.GPT_4O_MINI) + .addSystemMessage("You are a helpful assistant that can analyze images") + .addMessage(userMessage) + .temperature(0.0) + .build(); + + var response = openAIClient.chat().completions().create(request); + + // Verify the response + assertNotNull(response); + wireMock.verify(1, postRequestedFor(urlEqualTo("/chat/completions"))); + assertEquals("chatcmpl-test456", response.id()); + assertEquals( + "This image shows the Eiffel Tower in Paris, France.", + response.choices().get(0).message().content().get()); + + // Verify spans were exported + assertTrue( + openTelemetry + .getSdkTracerProvider() + .forceFlush() + .join(10, TimeUnit.SECONDS) + .isSuccess()); + var spanData = + getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow()); + assertNotNull(spanData); + assertEquals(1, spanData.size()); + var span = spanData.get(0); + + // Verify span attributes + assertEquals("openai", span.getAttributes().get(AttributeKey.stringKey("gen_ai.system"))); + assertEquals( + "gpt-4o-mini", + span.getAttributes().get(AttributeKey.stringKey("gen_ai.request.model"))); + assertEquals( + "gpt-4o-mini", + span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.model"))); + assertEquals( + "[stop]", + span.getAttributes() + .get(AttributeKey.stringArrayKey("gen_ai.response.finish_reasons")) + .toString()); + assertEquals( + "chat", span.getAttributes().get(AttributeKey.stringKey("gen_ai.operation.name"))); + assertEquals( + "chatcmpl-test456", + span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.id"))); + + // Verify input JSON captures both text and image content + String inputJson = + span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")); + assertNotNull(inputJson); + var inputMessages = JSON_MAPPER.readTree(inputJson); + assertEquals(2, inputMessages.size()); // system message + user message + + // Verify system message + var systemMessage = inputMessages.get(0); + assertEquals("system", systemMessage.get("role").asText()); + assertEquals( + "You are a helpful assistant that can analyze images", + systemMessage.get("content").asText()); + + // Verify user message with image + var userMsg = inputMessages.get(1); + assertEquals("user", userMsg.get("role").asText()); + assertTrue(userMsg.has("content")); + var content = userMsg.get("content"); + assertTrue(content.isArray()); + assertEquals(2, content.size()); // text + image + + // Verify text content part + var textContent = content.get(0); + assertEquals("text", textContent.get("type").asText()); + assertEquals("What's in this image?", textContent.get("text").asText()); + + // Verify image content part (now serialized as base64_attachment) + var imageContent = content.get(1); + assertEquals("base64_attachment", imageContent.get("type").asText()); + assertTrue(imageContent.has("content")); + assertTrue(imageContent.get("content").asText().startsWith("data:image/jpeg;base64,")); + assertEquals( + imageDataUrl, + imageContent.get("content").asText(), + "base64 data not correctly serialized"); + + // Verify usage metrics + assertEquals( + 150L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.input_tokens"))); + assertEquals( + 15L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.output_tokens"))); + + // Verify output JSON + String outputJson = + span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); + assertNotNull(outputJson); + var outputMessages = JSON_MAPPER.readTree(outputJson); + assertEquals(1, outputMessages.size()); + var messageZero = outputMessages.get(0); + assertEquals( + "This image shows the Eiffel Tower in Paris, France.", + messageZero.get("content").asText()); + } } diff --git a/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg b/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg new file mode 100644 index 0000000..f016466 Binary files /dev/null and b/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg differ diff --git a/src/test/java/dev/braintrust/trace/Base64AttachmentTest.java b/src/test/java/dev/braintrust/trace/Base64AttachmentTest.java new file mode 100644 index 0000000..bf3a7ce --- /dev/null +++ b/src/test/java/dev/braintrust/trace/Base64AttachmentTest.java @@ -0,0 +1,128 @@ +package dev.braintrust.trace; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.module.SimpleModule; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Base64; +import lombok.SneakyThrows; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class Base64AttachmentTest { + private static final ObjectMapper JSON_MAPPER = createObjectMapper(); + + private static ObjectMapper createObjectMapper() { + ObjectMapper mapper = new ObjectMapper(); + SimpleModule module = new SimpleModule(); + module.addSerializer(Base64Attachment.class, Base64Attachment.createSerializer()); + mapper.registerModule(module); + return mapper; + } + + @Test + void testOfWithValidDataUrl() { + String validDataUrl = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA"; + Base64Attachment attachment = Base64Attachment.of(validDataUrl); + assertNotNull(attachment); + } + + @Test + void testBadBase64Data() { + assertThrows(Exception.class, () -> Base64Attachment.of(null)); + assertThrows(Exception.class, () -> Base64Attachment.of("")); + } + + @Test + void testOfWithoutDataPrefixThrowsException() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> Base64Attachment.of("image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAUA")); + assertTrue(exception.getMessage().contains("data URL with format")); + } + + @Test + void testOfBase64WithoutMarkerThrowsException() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> Base64Attachment.of("data:image/png,iVBORw0KGgoAAAANSUhEUgAAAAUA")); + assertTrue(exception.getMessage().contains("data URL with format")); + } + + @Test + void testOfFileWithNonExistentFileThrowsException() { + RuntimeException exception = + assertThrows( + RuntimeException.class, + () -> + Base64Attachment.ofFile( + Base64Attachment.ContentType.IMAGE_PNG, + "/nonexistent/path/to/file.png")); + assertTrue(exception.getMessage().contains("Failed to read file")); + } + + @Test + @SneakyThrows + void testFileCreatesBase64Content(@TempDir Path tempDir) { + // Create a test file + Path testFile = tempDir.resolve("test.jpg"); + byte[] testData = "test jpeg data".getBytes(); + Files.write(testFile, testData); + + // Create attachment from file + Base64Attachment attachment = + Base64Attachment.ofFile( + Base64Attachment.ContentType.IMAGE_JPEG, testFile.toString()); + + // Serialize to JSON to verify the data URL format + String json = JSON_MAPPER.writeValueAsString(attachment); + + // Parse JSON and verify structure + var jsonNode = JSON_MAPPER.readTree(json); + assertEquals(2, jsonNode.size()); + + assertEquals("base64_attachment", jsonNode.get("type").asText()); + + String content = jsonNode.get("content").asText(); + assertTrue(content.startsWith("data:image/jpeg;base64,")); + + // Verify the base64 data is correct + String base64Part = content.substring("data:image/jpeg;base64,".length()); + byte[] decodedData = Base64.getDecoder().decode(base64Part); + assertArrayEquals(testData, decodedData); + } + + @Test + void testContentTypeOfNormalizesToLowercase() { + var customType = Base64Attachment.ContentType.of("IMAGE/PNG"); + assertEquals("image/png", customType.getMimeType()); + } + + @Test + void testContentTypeOfWithNullThrowsException() { + assertThrows(NullPointerException.class, () -> Base64Attachment.ContentType.of(null)); + } + + @Test + void testContentTypeEquality() { + var type1 = Base64Attachment.ContentType.of("image/png"); + var type2 = Base64Attachment.ContentType.of("image/png"); + var type3 = Base64Attachment.ContentType.of("image/jpeg"); + + assertEquals(type1, type2); + assertNotEquals(type1, type3); + assertEquals(type1, Base64Attachment.ContentType.IMAGE_PNG); + } + + @Test + void testContentTypeHashCode() { + var type1 = Base64Attachment.ContentType.of("image/png"); + var type2 = Base64Attachment.ContentType.of("image/png"); + + assertEquals(type1.hashCode(), type2.hashCode()); + } +}