Skip to content

Commit 4d36183

Browse files
realarkrealark
authored
Merge pull request #99 from braintrustdata/ark/oai-inst-streaming-bug
fix tagging for OpenAI response streaming
2 parents bb21361 + 7f762d8 commit 4d36183

34 files changed

Lines changed: 1498 additions & 440 deletions

File tree

braintrust-sdk/instrumentation/openai_2_8_0/build.gradle renamed to braintrust-sdk/instrumentation/openai_2_15_0/build.gradle

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ muzzle {
22
pass {
33
group = 'com.openai'
44
module = 'openai-java'
5-
versions = '[2.8.0,)'
5+
versions = '[2.15.0,)'
66
}
77
}
88

@@ -17,7 +17,7 @@ dependencies {
1717
compileOnly 'net.bytebuddy:byte-buddy:1.17.5'
1818

1919
// Target library — compileOnly because it will be on the app classpath at runtime
20-
compileOnly 'com.openai:openai-java:2.8.0'
20+
compileOnly 'com.openai:openai-java:2.15.0'
2121

2222
// Test dependencies
2323
testImplementation(testFixtures(project(":test-harness")))
@@ -26,7 +26,7 @@ dependencies {
2626
testRuntimeOnly 'org.junit.platform:junit-platform-launcher'
2727
testImplementation 'net.bytebuddy:byte-buddy-agent:1.17.5'
2828
testRuntimeOnly "org.slf4j:slf4j-simple:${slf4jVersion}"
29-
testImplementation 'com.openai:openai-java:2.8.0'
29+
testImplementation 'com.openai:openai-java:2.15.0'
3030
}
3131

3232
test {

braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java renamed to braintrust-sdk/instrumentation/openai_2_15_0/src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,13 @@
99
public class BraintrustOpenAI {
1010
/** Instrument openai client with braintrust traces */
1111
public static OpenAIClient wrapOpenAI(OpenTelemetry openTelemetry, OpenAIClient openAIClient) {
12-
return dev.braintrust.instrumentation.openai.v2_8_0.BraintrustOpenAI.wrapOpenAI(
12+
return dev.braintrust.instrumentation.openai.v2_15_0.BraintrustOpenAI.wrapOpenAI(
1313
openTelemetry, openAIClient);
1414
}
1515

1616
public static ChatCompletionCreateParams buildChatCompletionsPrompt(
1717
BraintrustPrompt prompt, Map<String, Object> parameters) {
18-
return dev.braintrust.instrumentation.openai.v2_8_0.BraintrustOpenAI
18+
return dev.braintrust.instrumentation.openai.v2_15_0.BraintrustOpenAI
1919
.buildChatCompletionsPrompt(prompt, parameters);
2020
}
2121
}

braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/BraintrustOpenAI.java renamed to braintrust-sdk/instrumentation/openai_2_15_0/src/main/java/dev/braintrust/instrumentation/openai/v2_15_0/BraintrustOpenAI.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package dev.braintrust.instrumentation.openai.v2_8_0;
1+
package dev.braintrust.instrumentation.openai.v2_15_0;
22

33
import com.openai.client.OpenAIClient;
44
import com.openai.core.ClientOptions;

braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/TracingHttpClient.java renamed to braintrust-sdk/instrumentation/openai_2_15_0/src/main/java/dev/braintrust/instrumentation/openai/v2_15_0/TracingHttpClient.java

Lines changed: 59 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
1-
package dev.braintrust.instrumentation.openai.v2_8_0;
1+
package dev.braintrust.instrumentation.openai.v2_15_0;
22

3+
import com.fasterxml.jackson.databind.json.JsonMapper;
4+
import com.openai.core.ObjectMappers;
35
import com.openai.core.RequestOptions;
4-
import com.openai.core.http.HttpClient;
5-
import com.openai.core.http.HttpRequest;
6-
import com.openai.core.http.HttpRequestBody;
7-
import com.openai.core.http.HttpResponse;
6+
import com.openai.core.http.*;
87
import com.openai.helpers.ChatCompletionAccumulator;
8+
import com.openai.helpers.ResponseAccumulator;
99
import com.openai.models.chat.completions.ChatCompletionChunk;
10+
import com.openai.models.responses.ResponseStreamEvent;
1011
import dev.braintrust.bootstrap.BraintrustBridge;
1112
import dev.braintrust.instrumentation.InstrumentationSemConv;
1213
import dev.braintrust.json.BraintrustJsonMapper;
@@ -23,6 +24,7 @@
2324

2425
@Slf4j
2526
class TracingHttpClient implements HttpClient {
27+
private static final JsonMapper JSON_MAPPER = ObjectMappers.jsonMapper();
2628
private final Tracer tracer;
2729
private final HttpClient underlying;
2830

@@ -173,7 +175,8 @@ private static void tagSpanFromBuffer(Span span, byte[] bytes, Long timeToFirstT
173175
if (bytes.length == 0) return;
174176
try {
175177
String firstLine = firstNonEmptyLine(bytes);
176-
if (firstLine != null && firstLine.startsWith("data:")) {
178+
if (firstLine != null
179+
&& (firstLine.startsWith("data:") || firstLine.startsWith("event:"))) {
177180
tagSpanFromSseBytes(span, bytes, timeToFirstTokenNanos);
178181
} else {
179182
InstrumentationSemConv.tagLLMSpanResponse(
@@ -205,26 +208,60 @@ private static String firstNonEmptyLine(byte[] bytes) {
205208
private static void tagSpanFromSseBytes(
206209
Span span, byte[] sseBytes, Long timeToFirstTokenNanos) {
207210
try {
208-
var accumulator = ChatCompletionAccumulator.create();
209211
var reader =
210212
new BufferedReader(
211213
new InputStreamReader(
212214
new ByteArrayInputStream(sseBytes), StandardCharsets.UTF_8));
213215
String line;
216+
String responseJson = null;
214217
while ((line = reader.readLine()) != null) {
215218
if (!line.startsWith("data:")) continue;
216-
String data = line.substring("data:".length()).strip();
217-
if (data.isEmpty() || data.equals("[DONE]")) continue;
218-
ChatCompletionChunk chunk =
219-
BraintrustJsonMapper.get().readValue(data, ChatCompletionChunk.class);
220-
accumulator.accumulate(chunk);
219+
var firstEventJson = line.substring("data:".length()).strip();
220+
// after the first data chunk is found, read the rest of the stream with the proper
221+
// accumulator type
222+
var jsonTree = JSON_MAPPER.readTree(firstEventJson);
223+
if (jsonTree.has("type") && jsonTree.get("type").asText().startsWith("response")) {
224+
// response API SSEvents
225+
ResponseAccumulator accumulator = ResponseAccumulator.create();
226+
accumulator.accumulate(
227+
JSON_MAPPER.readValue(firstEventJson, ResponseStreamEvent.class));
228+
while ((line = reader.readLine()) != null) {
229+
if (!line.startsWith("data:")) continue;
230+
String data = line.substring("data:".length()).strip();
231+
if (data.isEmpty() || data.equals("[DONE]")) continue;
232+
ResponseStreamEvent rse =
233+
JSON_MAPPER.readValue(data, ResponseStreamEvent.class);
234+
accumulator.accumulate(rse);
235+
}
236+
responseJson = JSON_MAPPER.writeValueAsString(accumulator.response());
237+
} else if (jsonTree.has("object")
238+
&& jsonTree.get("object").asText().equals("chat.completion.chunk")) {
239+
// completions API SSEvents
240+
var accumulator = ChatCompletionAccumulator.create();
241+
accumulator.accumulate(
242+
JSON_MAPPER.readValue(firstEventJson, ChatCompletionChunk.class));
243+
while ((line = reader.readLine()) != null) {
244+
if (!line.startsWith("data:")) continue;
245+
String data = line.substring("data:".length()).strip();
246+
if (data.isEmpty() || data.equals("[DONE]")) continue;
247+
ChatCompletionChunk chunk =
248+
BraintrustJsonMapper.get()
249+
.readValue(data, ChatCompletionChunk.class);
250+
accumulator.accumulate(chunk);
251+
}
252+
responseJson = JSON_MAPPER.writeValueAsString(accumulator.chatCompletion());
253+
} else {
254+
log.warn("unknown SSE object {}", firstEventJson);
255+
}
256+
break;
257+
}
258+
if (null != responseJson) {
259+
InstrumentationSemConv.tagLLMSpanResponse(
260+
span,
261+
InstrumentationSemConv.PROVIDER_NAME_OPENAI,
262+
responseJson,
263+
timeToFirstTokenNanos);
221264
}
222-
var chatCompletion = accumulator.chatCompletion();
223-
InstrumentationSemConv.tagLLMSpanResponse(
224-
span,
225-
InstrumentationSemConv.PROVIDER_NAME_OPENAI,
226-
BraintrustJsonMapper.toJson(chatCompletion),
227-
timeToFirstTokenNanos);
228265
} catch (Exception e) {
229266
log.error("Could not parse SSE buffer to tag streaming span output", e);
230267
}
@@ -233,8 +270,7 @@ private static void tagSpanFromSseBytes(
233270
/**
234271
* {@link HttpResponse} wrapper for streaming (SSE) responses. Its {@link #body()} returns a tee
235272
* {@link InputStream} that copies every byte the caller reads into an in-memory buffer. When
236-
* the stream is fully consumed and {@link #close()} is called, the accumulated bytes are
237-
* available via {@link #collectedBytes()} for span tagging.
273+
* the stream is fully consumed and {@link #close()} is called.
238274
*/
239275
private static final class TeeingStreamHttpResponse implements HttpResponse {
240276
private final HttpResponse delegate;
@@ -271,17 +307,13 @@ private void onStreamClosed() {
271307
}
272308
}
273309

274-
byte[] collectedBytes() {
275-
return teeBuffer.toByteArray();
276-
}
277-
278310
@Override
279311
public int statusCode() {
280312
return delegate.statusCode();
281313
}
282314

283315
@Override
284-
public com.openai.core.http.Headers headers() {
316+
public Headers headers() {
285317
return delegate.headers();
286318
}
287319

@@ -295,7 +327,7 @@ public void close() {
295327
try {
296328
teeStream.close(); // triggers onStreamClosed if not already fired (e.g. abandoned
297329
// stream)
298-
} catch (java.io.IOException ignored) {
330+
} catch (IOException ignored) {
299331
}
300332
delegate.close();
301333
}
@@ -322,7 +354,7 @@ private static final class TeeInputStream extends InputStream {
322354
}
323355

324356
@Override
325-
public int read() throws java.io.IOException {
357+
public int read() throws IOException {
326358
int b = source.read();
327359
if (b == -1) {
328360
notifyClosed();

braintrust-sdk/instrumentation/openai_2_8_0/src/main/java/dev/braintrust/instrumentation/openai/v2_8_0/auto/OpenAIInstrumentationModule.java renamed to braintrust-sdk/instrumentation/openai_2_15_0/src/main/java/dev/braintrust/instrumentation/openai/v2_15_0/auto/OpenAIInstrumentationModule.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
package dev.braintrust.instrumentation.openai.v2_8_0.auto;
1+
package dev.braintrust.instrumentation.openai.v2_15_0.auto;
22

33
import static net.bytebuddy.matcher.ElementMatchers.named;
44
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
@@ -8,7 +8,7 @@
88
import dev.braintrust.instrumentation.InstrumentationModule;
99
import dev.braintrust.instrumentation.TypeInstrumentation;
1010
import dev.braintrust.instrumentation.TypeTransformer;
11-
import dev.braintrust.instrumentation.openai.v2_8_0.BraintrustOpenAI;
11+
import dev.braintrust.instrumentation.openai.v2_15_0.BraintrustOpenAI;
1212
import io.opentelemetry.api.GlobalOpenTelemetry;
1313
import java.util.List;
1414
import java.util.Set;
@@ -20,10 +20,10 @@
2020
@AutoService(InstrumentationModule.class)
2121
public class OpenAIInstrumentationModule extends InstrumentationModule {
2222
private static final String MANUAL_INSTRUMENTATION_PACKAGE =
23-
"dev.braintrust.instrumentation.openai.v2_8_0.";
23+
"dev.braintrust.instrumentation.openai.v2_15_0.";
2424

2525
public OpenAIInstrumentationModule() {
26-
super("openai_2_8_0");
26+
super("openai_2_15_0");
2727
}
2828

2929
@Override

0 commit comments

Comments
 (0)