Skip to content

Commit 989a1bf

Browse files
realarkrealark
committed
wip
1 parent 73f2b5f commit 989a1bf

4 files changed

Lines changed: 322 additions & 5 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ private static void chatCompletionsExample(OpenAIClient openAIClient) {
4949
String imageDataUrl =
5050
Base64Attachment.ofFile(
5151
Base64Attachment.ContentType.IMAGE_JPEG,
52-
"/home/ark/braintrust/java/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg")
52+
"/home/ark/braintrust/sdk-java/src/test/java/dev/braintrust/instrumentation/openai/travel-paris-france-poster.jpg")
5353
.getBase64Data();
5454

5555
// Create text content part

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

Lines changed: 317 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,16 @@
1212
import com.fasterxml.jackson.databind.ObjectMapper;
1313
import com.fasterxml.jackson.databind.SerializerProvider;
1414
import com.fasterxml.jackson.databind.module.SimpleModule;
15+
import com.fasterxml.jackson.annotation.JsonInclude;
16+
import com.fasterxml.jackson.annotation.JsonProperty;
1517
import com.openai.models.chat.completions.ChatCompletion;
1618
import com.openai.models.chat.completions.ChatCompletionAssistantMessageParam;
1719
import com.openai.models.chat.completions.ChatCompletionContentPartImage;
1820
import com.openai.models.chat.completions.ChatCompletionContentPartText;
1921
import com.openai.models.chat.completions.ChatCompletionCreateParams;
2022
import com.openai.models.chat.completions.ChatCompletionDeveloperMessageParam;
2123
import com.openai.models.chat.completions.ChatCompletionMessage;
24+
import com.openai.models.chat.completions.ChatCompletionMessageParam;
2225
import com.openai.models.chat.completions.ChatCompletionMessageToolCall;
2326
import com.openai.models.chat.completions.ChatCompletionSystemMessageParam;
2427
import com.openai.models.chat.completions.ChatCompletionToolMessageParam;
@@ -34,6 +37,7 @@
3437
import java.lang.invoke.MethodHandle;
3538
import java.lang.invoke.MethodHandles;
3639
import java.lang.invoke.MethodType;
40+
import java.util.ArrayList;
3741
import java.util.HashMap;
3842
import java.util.List;
3943
import java.util.Map;
@@ -54,6 +58,7 @@ private static ObjectMapper createObjectMapper() {
5458
final JsonSerializer<Base64Attachment> attachmentSerializer =
5559
Base64Attachment.createSerializer();
5660
ObjectMapper mapper = new ObjectMapper();
61+
mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
5762
SimpleModule module = new SimpleModule();
5863
module.addSerializer(
5964
ChatCompletionContentPartImage.class,
@@ -81,16 +86,324 @@ public void serialize(
8186
return mapper;
8287
}
8388

89+
// OTel GenAI Semantic Convention structures
90+
static class SemconvChatMessage {
91+
@JsonProperty("role")
92+
public final String role;
93+
94+
@JsonProperty("parts")
95+
public final List<Object> parts;
96+
97+
public SemconvChatMessage(String role, List<Object> parts) {
98+
this.role = role;
99+
this.parts = parts;
100+
}
101+
}
102+
103+
@JsonInclude(JsonInclude.Include.NON_NULL)
104+
static class TextPart {
105+
@JsonProperty("type")
106+
public final String type = "text";
107+
108+
@JsonProperty("content")
109+
public final String content;
110+
111+
public TextPart(String content) {
112+
this.content = content;
113+
}
114+
}
115+
116+
@JsonInclude(JsonInclude.Include.NON_NULL)
117+
static class ToolCallRequestPart {
118+
@JsonProperty("type")
119+
public final String type = "tool_call";
120+
121+
@JsonProperty("name")
122+
public final String name;
123+
124+
@JsonProperty("id")
125+
public final String id;
126+
127+
@JsonProperty("arguments")
128+
public final String arguments;
129+
130+
public ToolCallRequestPart(String name, String id, String arguments) {
131+
this.name = name;
132+
this.id = id;
133+
this.arguments = arguments;
134+
}
135+
}
136+
137+
@JsonInclude(JsonInclude.Include.NON_NULL)
138+
static class ToolCallResponsePart {
139+
@JsonProperty("type")
140+
public final String type = "tool_call_response";
141+
142+
@JsonProperty("id")
143+
public final String id;
144+
145+
@JsonProperty("response")
146+
public final String response;
147+
148+
public ToolCallResponsePart(String id, String response) {
149+
this.id = id;
150+
this.response = response;
151+
}
152+
}
153+
154+
@JsonInclude(JsonInclude.Include.NON_NULL)
155+
static class GenericPart {
156+
@JsonProperty("type")
157+
public final String type;
158+
159+
@JsonProperty("data")
160+
public final Object data;
161+
162+
public GenericPart(String type, Object data) {
163+
this.type = type;
164+
this.data = data;
165+
}
166+
}
167+
168+
// Transform OpenAI messages to OTel GenAI semconv format
169+
private static List<SemconvChatMessage> transformToSemconvMessages(
170+
List<ChatCompletionMessageParam> messages) {
171+
return messages.stream()
172+
.map(ChatCompletionEventsHelper::transformMessage)
173+
.filter(Objects::nonNull)
174+
.collect(Collectors.toList());
175+
}
176+
177+
@Nullable
178+
private static SemconvChatMessage transformMessage(ChatCompletionMessageParam message) {
179+
// Try each union type accessor
180+
var userOpt = tryGetUser(message);
181+
if (userOpt.isPresent()) {
182+
return transformUserMessage(userOpt.get());
183+
}
184+
185+
var systemOpt = tryGetSystem(message);
186+
if (systemOpt.isPresent()) {
187+
return transformSystemMessage(systemOpt.get());
188+
}
189+
190+
var assistantOpt = tryGetAssistant(message);
191+
if (assistantOpt.isPresent()) {
192+
return transformAssistantMessage(assistantOpt.get());
193+
}
194+
195+
var toolOpt = tryGetTool(message);
196+
if (toolOpt.isPresent()) {
197+
return transformToolMessage(toolOpt.get());
198+
}
199+
200+
var developerOpt = tryGetDeveloper(message);
201+
if (developerOpt.isPresent()) {
202+
return transformDeveloperMessage(developerOpt.get());
203+
}
204+
205+
return null;
206+
}
207+
208+
private static Optional<ChatCompletionUserMessageParam> tryGetUser(
209+
ChatCompletionMessageParam message) {
210+
try {
211+
var method = message.getClass().getMethod("user");
212+
@SuppressWarnings("unchecked")
213+
var result = (Optional<ChatCompletionUserMessageParam>) method.invoke(message);
214+
return result;
215+
} catch (Exception e) {
216+
return Optional.empty();
217+
}
218+
}
219+
220+
private static Optional<ChatCompletionSystemMessageParam> tryGetSystem(
221+
ChatCompletionMessageParam message) {
222+
try {
223+
var method = message.getClass().getMethod("system");
224+
@SuppressWarnings("unchecked")
225+
var result = (Optional<ChatCompletionSystemMessageParam>) method.invoke(message);
226+
return result;
227+
} catch (Exception e) {
228+
return Optional.empty();
229+
}
230+
}
231+
232+
private static Optional<ChatCompletionAssistantMessageParam> tryGetAssistant(
233+
ChatCompletionMessageParam message) {
234+
try {
235+
var method = message.getClass().getMethod("assistant");
236+
@SuppressWarnings("unchecked")
237+
var result = (Optional<ChatCompletionAssistantMessageParam>) method.invoke(message);
238+
return result;
239+
} catch (Exception e) {
240+
return Optional.empty();
241+
}
242+
}
243+
244+
private static Optional<ChatCompletionToolMessageParam> tryGetTool(
245+
ChatCompletionMessageParam message) {
246+
try {
247+
var method = message.getClass().getMethod("tool");
248+
@SuppressWarnings("unchecked")
249+
var result = (Optional<ChatCompletionToolMessageParam>) method.invoke(message);
250+
return result;
251+
} catch (Exception e) {
252+
return Optional.empty();
253+
}
254+
}
255+
256+
private static Optional<ChatCompletionDeveloperMessageParam> tryGetDeveloper(
257+
ChatCompletionMessageParam message) {
258+
try {
259+
var method = message.getClass().getMethod("developer");
260+
@SuppressWarnings("unchecked")
261+
var result = (Optional<ChatCompletionDeveloperMessageParam>) method.invoke(message);
262+
return result;
263+
} catch (Exception e) {
264+
return Optional.empty();
265+
}
266+
}
267+
268+
private static SemconvChatMessage transformUserMessage(ChatCompletionUserMessageParam message) {
269+
List<Object> parts = new ArrayList<>();
270+
ChatCompletionUserMessageParam.Content content = message.content();
271+
272+
if (content.isText()) {
273+
parts.add(new TextPart(content.asText()));
274+
} else if (content.isArrayOfContentParts()) {
275+
for (var part : content.asArrayOfContentParts()) {
276+
if (part.isText()) {
277+
parts.add(new TextPart(part.asText().text()));
278+
} else {
279+
// Try to get image part using union type accessor
280+
var imageOpt = tryGetImagePart(part);
281+
if (imageOpt.isPresent()) {
282+
try {
283+
var imageUrl = imageOpt.get().imageUrl().url();
284+
var attachment = Base64Attachment.of(imageUrl);
285+
parts.add(attachment);
286+
} catch (Exception e) {
287+
log.debug("Failed to parse image URL", e);
288+
}
289+
}
290+
}
291+
}
292+
}
293+
294+
return new SemconvChatMessage("user", parts);
295+
}
296+
297+
private static Optional<ChatCompletionContentPartImage> tryGetImagePart(
298+
com.openai.models.chat.completions.ChatCompletionContentPart part) {
299+
try {
300+
var method = part.getClass().getMethod("imageUrl");
301+
@SuppressWarnings("unchecked")
302+
var result = (Optional<ChatCompletionContentPartImage>) method.invoke(part);
303+
return result;
304+
} catch (Exception e) {
305+
return Optional.empty();
306+
}
307+
}
308+
309+
private static SemconvChatMessage transformSystemMessage(
310+
ChatCompletionSystemMessageParam message) {
311+
List<Object> parts = new ArrayList<>();
312+
ChatCompletionSystemMessageParam.Content content = message.content();
313+
314+
if (content.isText()) {
315+
parts.add(new TextPart(content.asText()));
316+
} else if (content.isArrayOfContentParts()) {
317+
for (var part : content.asArrayOfContentParts()) {
318+
parts.add(new TextPart(part.text()));
319+
}
320+
}
321+
322+
return new SemconvChatMessage("system", parts);
323+
}
324+
325+
private static SemconvChatMessage transformDeveloperMessage(
326+
ChatCompletionDeveloperMessageParam message) {
327+
List<Object> parts = new ArrayList<>();
328+
ChatCompletionDeveloperMessageParam.Content content = message.content();
329+
330+
if (content.isText()) {
331+
parts.add(new TextPart(content.asText()));
332+
} else if (content.isArrayOfContentParts()) {
333+
for (var part : content.asArrayOfContentParts()) {
334+
parts.add(new TextPart(part.text()));
335+
}
336+
}
337+
338+
return new SemconvChatMessage("system", parts);
339+
}
340+
341+
private static SemconvChatMessage transformAssistantMessage(
342+
ChatCompletionAssistantMessageParam message) {
343+
List<Object> parts = new ArrayList<>();
344+
345+
// Handle text content
346+
message.content()
347+
.ifPresent(
348+
content -> {
349+
if (content.isText()) {
350+
parts.add(new TextPart(content.asText()));
351+
} else if (content.isArrayOfContentParts()) {
352+
for (var part : content.asArrayOfContentParts()) {
353+
if (part.isText()) {
354+
parts.add(new TextPart(part.asText().text()));
355+
} else if (part.isRefusal()) {
356+
parts.add(new TextPart(part.asRefusal().refusal()));
357+
}
358+
}
359+
}
360+
});
361+
362+
// Handle tool calls
363+
message.toolCalls()
364+
.ifPresent(
365+
toolCalls -> {
366+
for (var toolCall : toolCalls) {
367+
FunctionAccess functionAccess = getFunctionAccess(toolCall);
368+
if (functionAccess != null) {
369+
parts.add(
370+
new ToolCallRequestPart(
371+
functionAccess.name(),
372+
functionAccess.id(),
373+
functionAccess.arguments()));
374+
}
375+
}
376+
});
377+
378+
return new SemconvChatMessage("assistant", parts);
379+
}
380+
381+
private static SemconvChatMessage transformToolMessage(ChatCompletionToolMessageParam message) {
382+
List<Object> parts = new ArrayList<>();
383+
String toolCallId = message.toolCallId();
384+
ChatCompletionToolMessageParam.Content content = message.content();
385+
386+
String responseContent = "";
387+
if (content.isText()) {
388+
responseContent = content.asText();
389+
} else if (content.isArrayOfContentParts()) {
390+
responseContent = joinContentParts(content.asArrayOfContentParts());
391+
}
392+
393+
parts.add(new ToolCallResponsePart(toolCallId, responseContent));
394+
return new SemconvChatMessage("tool", parts);
395+
}
396+
84397
@SneakyThrows
85398
public static void emitPromptLogEvents(
86399
Context context,
87400
Logger eventLogger,
88401
ChatCompletionCreateParams request,
89402
boolean captureMessageContent) {
90-
Span.current()
91-
.setAttribute(
92-
"gen_ai.input.messages",
93-
JSON_MAPPER.writeValueAsString(request.messages()));
403+
List<SemconvChatMessage> semconvMessages = transformToSemconvMessages(request.messages());
404+
String semconvJson = JSON_MAPPER.writeValueAsString(semconvMessages);
405+
Span span = Span.current();
406+
span.setAttribute("gen_ai.input.messages", semconvJson);
94407
}
95408

96409
private static String contentToString(ChatCompletionToolMessageParam.Content content) {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package dev.braintrust.trace;
22

3+
import com.fasterxml.jackson.annotation.JsonProperty;
34
import com.fasterxml.jackson.core.JsonGenerator;
45
import com.fasterxml.jackson.databind.JsonSerializer;
56
import com.fasterxml.jackson.databind.SerializerProvider;
@@ -23,7 +24,9 @@
2324
* https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-output-messages.json
2425
*/
2526
public class Base64Attachment {
27+
@JsonProperty("type")
2628
@Getter private final String type = "base64_attachment";
29+
@JsonProperty("content")
2730
@Getter private final String base64Data;
2831

2932
private Base64Attachment(@Nonnull String base64Data) {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ public CompletableResultCode export(Collection<SpanData> spans) {
4848

4949
// Combine all results
5050
var combined = CompletableResultCode.ofAll(results);
51+
log.debug("span export results: {}", combined.isSuccess());
5152

5253
return combined;
5354
}

0 commit comments

Comments
 (0)