Skip to content

Commit 9b52bba

Browse files
realarkrealark
authored
Merge pull request #22 from braintrustdata/ark/oai-and-jar-tests
OpenAI instrumentation test
2 parents ca9f999 + 0b542f8 commit 9b52bba

2 files changed

Lines changed: 173 additions & 0 deletions

File tree

‎build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ dependencies {
5252

5353
testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"
5454
testImplementation "org.junit.jupiter:junit-jupiter-params:${junitVersion}"
55+
testImplementation 'com.openai:openai-java:2.8.1'
56+
testImplementation 'com.github.tomakehurst:wiremock-jre8:2.35.0'
5557
}
5658

5759
test {

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

Lines changed: 171 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
package dev.braintrust.instrumentation.openai;
2+
3+
import static com.github.tomakehurst.wiremock.client.WireMock.*;
4+
import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
5+
import static dev.braintrust.trace.BraintrustTracingTest.getExportedBraintrustSpans;
6+
import static org.junit.jupiter.api.Assertions.*;
7+
8+
import com.fasterxml.jackson.databind.ObjectMapper;
9+
import com.github.tomakehurst.wiremock.junit5.WireMockExtension;
10+
import com.openai.client.OpenAIClient;
11+
import com.openai.client.okhttp.OpenAIOkHttpClient;
12+
import com.openai.models.ChatModel;
13+
import com.openai.models.chat.completions.ChatCompletionCreateParams;
14+
import dev.braintrust.config.BraintrustConfig;
15+
import dev.braintrust.trace.BraintrustTracing;
16+
import io.opentelemetry.api.GlobalOpenTelemetry;
17+
import io.opentelemetry.api.common.AttributeKey;
18+
import io.opentelemetry.sdk.OpenTelemetrySdk;
19+
import java.util.concurrent.TimeUnit;
20+
import org.junit.jupiter.api.BeforeEach;
21+
import org.junit.jupiter.api.Test;
22+
import org.junit.jupiter.api.extension.RegisterExtension;
23+
24+
public class BraintrustOpenAITest {
25+
private static final ObjectMapper JSON_MAPPER =
26+
new com.fasterxml.jackson.databind.ObjectMapper();
27+
28+
@RegisterExtension
29+
static WireMockExtension wireMock =
30+
WireMockExtension.newInstance().options(wireMockConfig().dynamicPort()).build();
31+
32+
private final BraintrustConfig config =
33+
BraintrustConfig.of(
34+
"BRAINTRUST_API_KEY", "foobar",
35+
"BRAINTRUST_DEFAULT_PROJECT_NAME", "unit-test-project",
36+
"BRAINTRUST_JAVA_EXPORT_SPANS_IN_MEMORY_FOR_UNIT_TEST", "true");
37+
38+
@BeforeEach
39+
void beforeEach() {
40+
GlobalOpenTelemetry.resetForTest();
41+
getExportedBraintrustSpans().clear();
42+
wireMock.resetAll();
43+
}
44+
45+
@Test
46+
void testWrapOpenAi() {
47+
// Mock the OpenAI API response
48+
wireMock.stubFor(
49+
post(urlEqualTo("/chat/completions"))
50+
.willReturn(
51+
aResponse()
52+
.withStatus(200)
53+
.withHeader("Content-Type", "application/json")
54+
.withBody(
55+
"""
56+
{
57+
"id": "chatcmpl-test123",
58+
"object": "chat.completion",
59+
"created": 1677652288,
60+
"model": "gpt-4o-mini",
61+
"choices": [
62+
{
63+
"index": 0,
64+
"message": {
65+
"role": "assistant",
66+
"content": "The capital of France is Paris."
67+
},
68+
"finish_reason": "stop"
69+
}
70+
],
71+
"usage": {
72+
"prompt_tokens": 20,
73+
"completion_tokens": 8,
74+
"total_tokens": 28
75+
}
76+
}
77+
""")));
78+
79+
var openTelemetry = (OpenTelemetrySdk) BraintrustTracing.of(config, true);
80+
81+
// Create OpenAI client pointing to WireMock server
82+
OpenAIClient openAIClient =
83+
OpenAIOkHttpClient.builder()
84+
.baseUrl("http://localhost:" + wireMock.getPort())
85+
.apiKey("test-api-key")
86+
.build();
87+
88+
// Wrap with Braintrust instrumentation
89+
openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, openAIClient);
90+
91+
var request =
92+
ChatCompletionCreateParams.builder()
93+
.model(ChatModel.GPT_4O_MINI)
94+
.addSystemMessage("You are a helpful assistant")
95+
.addUserMessage("What is the capital of France?")
96+
.temperature(0.0)
97+
.build();
98+
99+
var response = openAIClient.chat().completions().create(request);
100+
101+
// Verify the response
102+
assertNotNull(response);
103+
wireMock.verify(1, postRequestedFor(urlEqualTo("/chat/completions")));
104+
assertEquals("chatcmpl-test123", response.id());
105+
assertEquals(
106+
"The capital of France is Paris.",
107+
response.choices().get(0).message().content().get());
108+
109+
// Verify spans were exported
110+
assertTrue(
111+
openTelemetry
112+
.getSdkTracerProvider()
113+
.forceFlush()
114+
.join(10, TimeUnit.SECONDS)
115+
.isSuccess());
116+
var spanData =
117+
getExportedBraintrustSpans().get(config.getBraintrustParentValue().orElseThrow());
118+
assertNotNull(spanData);
119+
assertEquals(1, spanData.size());
120+
var span = spanData.get(0);
121+
122+
assertEquals("openai", span.getAttributes().get(AttributeKey.stringKey("gen_ai.system")));
123+
assertEquals(
124+
"gpt-4o-mini",
125+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.request.model")));
126+
assertEquals(
127+
"gpt-4o-mini",
128+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.model")));
129+
assertEquals(
130+
"[stop]",
131+
span.getAttributes()
132+
.get(AttributeKey.stringArrayKey("gen_ai.response.finish_reasons"))
133+
.toString());
134+
assertEquals(
135+
"chat", span.getAttributes().get(AttributeKey.stringKey("gen_ai.operation.name")));
136+
assertEquals(
137+
"chatcmpl-test123",
138+
span.getAttributes().get(AttributeKey.stringKey("gen_ai.response.id")));
139+
assertEquals(
140+
"You are a helpful assistant",
141+
span.getAttributes().get(AttributeKey.stringKey("instructions")));
142+
assertEquals(
143+
"\"What is the capital of France?\"",
144+
span.getAttributes().get(AttributeKey.stringKey("braintrust.input_json")));
145+
assertEquals(
146+
"project_name:unit-test-project",
147+
span.getAttributes().get(AttributeKey.stringKey("braintrust.parent")));
148+
assertEquals(
149+
20L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.input_tokens")));
150+
assertEquals(
151+
8L, span.getAttributes().get(AttributeKey.longKey("gen_ai.usage.output_tokens")));
152+
assertEquals(
153+
0.0,
154+
span.getAttributes().get(AttributeKey.doubleKey("gen_ai.request.temperature")));
155+
String outputJson =
156+
span.getAttributes().get(AttributeKey.stringKey("braintrust.output_json"));
157+
assertNotNull(outputJson);
158+
try {
159+
var jsonNode = JSON_MAPPER.readTree(outputJson);
160+
assertEquals("chatcmpl-test123", jsonNode.get("id").asText());
161+
assertEquals(
162+
"The capital of France is Paris.",
163+
jsonNode.get("choices").get(0).get("message").get("content").asText());
164+
assertEquals(8, jsonNode.get("usage").get("completion_tokens").asInt());
165+
assertEquals(20, jsonNode.get("usage").get("prompt_tokens").asInt());
166+
assertEquals(28, jsonNode.get("usage").get("total_tokens").asInt());
167+
} catch (Exception e) {
168+
fail("Failed to parse output JSON: " + e.getMessage());
169+
}
170+
}
171+
}

0 commit comments

Comments
 (0)