From 0b542f853c266033c39046e9a5fc60a9bf856c8d Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Tue, 30 Sep 2025 17:07:14 -0600 Subject: [PATCH] OpenAI instrumentation test --- build.gradle | 2 + .../openai/BraintrustOpenAITest.java | 171 ++++++++++++++++++ 2 files changed, 173 insertions(+) create mode 100644 src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java diff --git a/build.gradle b/build.gradle index 7cd08cc..0beadef 100644 --- a/build.gradle +++ b/build.gradle @@ -52,6 +52,8 @@ dependencies { testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}" testImplementation "org.junit.jupiter:junit-jupiter-params:${junitVersion}" + testImplementation 'com.openai:openai-java:2.8.1' + testImplementation 'com.github.tomakehurst:wiremock-jre8:2.35.0' } test { diff --git a/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java b/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java new file mode 100644 index 0000000..6010106 --- /dev/null +++ b/src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java @@ -0,0 +1,171 @@ +package dev.braintrust.instrumentation.openai; + +import static com.github.tomakehurst.wiremock.client.WireMock.*; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; +import static dev.braintrust.trace.BraintrustTracingTest.getExportedBraintrustSpans; +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.github.tomakehurst.wiremock.junit5.WireMockExtension; +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.models.ChatModel; +import com.openai.models.chat.completions.ChatCompletionCreateParams; +import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.trace.BraintrustTracing; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +public class BraintrustOpenAITest { + private static final ObjectMapper JSON_MAPPER = + new com.fasterxml.jackson.databind.ObjectMapper(); + + @RegisterExtension + static WireMockExtension wireMock = + WireMockExtension.newInstance().options(wireMockConfig().dynamicPort()).build(); + + private final BraintrustConfig config = + BraintrustConfig.of( + "BRAINTRUST_API_KEY", "foobar", + "BRAINTRUST_DEFAULT_PROJECT_NAME", "unit-test-project", + "BRAINTRUST_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "true"); + + @BeforeEach + void beforeEach() { + GlobalOpenTelemetry.resetForTest(); + getExportedBraintrustSpans().clear(); + wireMock.resetAll(); + } + + @Test + void testWrapOpenAi() { + // Mock the OpenAI API response + wireMock.stubFor( + post(urlEqualTo("/chat/completions")) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "id": "chatcmpl-test123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The capital of France is Paris." + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 20, + "completion_tokens": 8, + "total_tokens": 28 + } + } + """))); + + 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); + + var request = + ChatCompletionCreateParams.builder() + .model(ChatModel.GPT_4O_MINI) + .addSystemMessage("You are a helpful assistant") + .addUserMessage("What is the capital of France?") + .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-test123", response.id()); + assertEquals( + "The capital of France is Paris.", + 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); + + 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-test123", + span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.id"))); + assertEquals( + "You are a helpful assistant", + span.getAttributes().get(AttributeKey.stringKey("instructions"))); + assertEquals( + "\"What is the capital of France?\"", + span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json"))); + assertEquals( + "project_name:unit-test-project", + span.getAttributes().get(AttributeKey.stringKey("braintrust.parent"))); + assertEquals( + 20L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.input_tokens"))); + assertEquals( + 8L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.output_tokens"))); + assertEquals( + 0.0, + span.getAttributes().get(AttributeKey.doubleKey("gen_ai.request.temperature"))); + String outputJson = + span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json")); + assertNotNull(outputJson); + try { + var jsonNode = JSON_MAPPER.readTree(outputJson); + assertEquals("chatcmpl-test123", jsonNode.get("id").asText()); + assertEquals( + "The capital of France is Paris.", + jsonNode.get("choices").get(0).get("message").get("content").asText()); + assertEquals(8, jsonNode.get("usage").get("completion_tokens").asInt()); + assertEquals(20, jsonNode.get("usage").get("prompt_tokens").asInt()); + assertEquals(28, jsonNode.get("usage").get("total_tokens").asInt()); + } catch (Exception e) { + fail("Failed to parse output JSON: " + e.getMessage()); + } + } +}