Skip to content

add option to configure braintrust SDK json object mapper #34

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
realark merged 1 commit into from
Feb 11, 2026
Merged

add option to configure braintrust SDK json object mapper #34

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 10 additions & 15 deletions src/main/java/com/google/genai/BraintrustApiClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package com.google.genai;

import com.fasterxml.jackson.databind.ObjectMapper;
import static dev.braintrust.json.BraintrustJsonMapper.fromJson;
import static dev.braintrust.json.BraintrustJsonMapper.toJson;

import com.google.genai.types.HttpOptions;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.api.trace.Span;
Expand All @@ -27,8 +29,6 @@
*/
@Slf4j
class BraintrustApiClient extends ApiClient {
private static final ObjectMapper JSON_MAPPER = new ObjectMapper();

private final ApiClient delegate;
private final Tracer tracer;

Expand Down Expand Up @@ -58,7 +58,7 @@ private void tagSpan(

// Parse request
if (requestBody != null) {
var requestJson = JSON_MAPPER.readValue(requestBody, Map.class);
var requestJson = fromJson(requestBody, Map.class);

// Extract metadata fields
for (String field :
Expand Down Expand Up @@ -108,22 +108,20 @@ private void tagSpan(
inputJson.put("config", requestJson.get("generationConfig"));
}

span.setAttribute(
"braintrust.input_json", JSON_MAPPER.writeValueAsString(inputJson));
span.setAttribute("braintrust.input_json", toJson(inputJson));
}

// Parse response
if (responseBody != null) {
var responseJson = JSON_MAPPER.readValue(responseBody, Map.class);
var responseJson = fromJson(responseBody, Map.class);

// Extract model version from response
if (responseJson.containsKey("modelVersion")) {
metadata.put("model", responseJson.get("modelVersion"));
}

// Set full response as output_json
span.setAttribute(
"braintrust.output_json", JSON_MAPPER.writeValueAsString(responseJson));
span.setAttribute("braintrust.output_json", toJson(responseJson));

// Parse usage metadata for metrics
if (responseJson.get("usageMetadata") instanceof Map) {
Expand All @@ -146,18 +144,15 @@ private void tagSpan(
(Number) usage.get("cachedContentTokenCount"));
}

span.setAttribute(
"braintrust.metrics", JSON_MAPPER.writeValueAsString(metrics));
span.setAttribute("braintrust.metrics", toJson(metrics));
}
}

// Set metadata
span.setAttribute("braintrust.metadata", JSON_MAPPER.writeValueAsString(metadata));
span.setAttribute("braintrust.metadata", toJson(metadata));

// Set span_attributes to mark as LLM span
span.setAttribute(
"braintrust.span_attributes",
JSON_MAPPER.writeValueAsString(Map.of("type", "llm")));
span.setAttribute("braintrust.span_attributes", toJson(Map.of("type", "llm")));

} catch (Throwable t) {
log.warn("failed to tag gemini span", t);
Expand Down
62 changes: 17 additions & 45 deletions src/main/java/dev/braintrust/api/BraintrustApiClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,9 @@
package dev.braintrust.api;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategies;
import com.fasterxml.jackson.datatype.jdk8.Jdk8Module;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import static dev.braintrust.json.BraintrustJsonMapper.fromJson;
import static dev.braintrust.json.BraintrustJsonMapper.toJson;

import dev.braintrust.config.BraintrustConfig;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
Expand Down Expand Up @@ -111,7 +108,6 @@ static BraintrustApiClient of(BraintrustConfig config) {
class HttpImpl implements BraintrustApiClient {
private final BraintrustConfig config;
private final HttpClient httpClient;
private final ObjectMapper objectMapper;

HttpImpl(BraintrustConfig config) {
this(config, createDefaultHttpClient(config));
Expand All @@ -120,7 +116,6 @@ class HttpImpl implements BraintrustApiClient {
private HttpImpl(BraintrustConfig config, HttpClient httpClient) {
this.config = config;
this.httpClient = httpClient;
this.objectMapper = createObjectMapper();
}

@Override
Expand Down Expand Up @@ -422,24 +417,19 @@ private <T> CompletableFuture<T> getAsync(String path, Class<T> responseType) {

private <T> CompletableFuture<T> postAsync(
String path, Object body, Class<T> responseType) {
try {
var jsonBody = objectMapper.writeValueAsString(body);

var request =
HttpRequest.newBuilder()
.uri(URI.create(config.apiUrl() + path))
.header("Authorization", "Bearer " + config.apiKey())
.header("Content-Type", "application/json")
.header("Accept", "application/json")
.timeout(config.requestTimeout())
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
.build();

return sendAsync(request, responseType);
} catch (IOException e) {
return CompletableFuture.failedFuture(
new ApiException("Failed to serialize request body", e));
}
var jsonBody = toJson(body);

var request =
HttpRequest.newBuilder()
.uri(URI.create(config.apiUrl() + path))
.header("Authorization", "Bearer " + config.apiKey())
.header("Content-Type", "application/json")
.header("Accept", "application/json")
.timeout(config.requestTimeout())
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
.build();

return sendAsync(request, responseType);
}

private <T> CompletableFuture<T> sendAsync(HttpRequest request, Class<T> responseType) {
Expand All @@ -454,12 +444,7 @@ private <T> T handleResponse(HttpResponse<String> response, Class<T> responseTyp
log.debug("API Response: {} - {}", response.statusCode(), response.body());

if (response.statusCode() >= 200 && response.statusCode() < 300) {
try {
return objectMapper.readValue(response.body(), responseType);
} catch (IOException e) {
log.warn("Failed to parse response body", e);
throw new ApiException("Failed to parse response body", e);
}
return fromJson(response.body(), responseType);
} else {
log.warn(
"API request failed with status {}: {}",
Expand Down Expand Up @@ -488,19 +473,6 @@ private boolean isNotFound(Throwable error) {
private static HttpClient createDefaultHttpClient(BraintrustConfig config) {
return HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build();
}

private static ObjectMapper createObjectMapper() {
return new ObjectMapper()
.registerModule(new JavaTimeModule())
.registerModule(new Jdk8Module()) // For Optional support
.setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE)
.setSerializationInclusion(
JsonInclude.Include.NON_ABSENT) // Skip null and absent Optional
.configure(
com.fasterxml.jackson.databind.DeserializationFeature
.FAIL_ON_UNKNOWN_PROPERTIES,
false); // Ignore unknown fields from API
}
}

/** Implementation for test doubling */
Expand Down
52 changes: 20 additions & 32 deletions src/main/java/dev/braintrust/devserver/Devserver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
package dev.braintrust.devserver;

import com.fasterxml.jackson.databind.ObjectMapper;
import static dev.braintrust.json.BraintrustJsonMapper.fromJson;
import static dev.braintrust.json.BraintrustJsonMapper.toJson;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpServer;
Expand Down Expand Up @@ -83,11 +85,6 @@ public class Devserver {
private final @Nullable Consumer<io.opentelemetry.sdk.trace.SdkTracerProviderBuilder>
traceBuilderHook;
private final @Nullable Consumer<BraintrustConfig.Builder> configBuilderHook;
private static final ObjectMapper JSON_MAPPER =
new ObjectMapper()
.enable(
com.fasterxml.jackson.core.JsonParser.Feature
.INCLUDE_SOURCE_IN_LOCATION);

// LRU cache for token -> Braintrust mappings
private final LRUCache<String, Braintrust> authCache = new LRUCache<>(32);
Expand Down Expand Up @@ -219,7 +216,7 @@ private void handleList(HttpExchange exchange) throws IOException {
response.put(evalName, metadata);
}

String jsonResponse = JSON_MAPPER.writeValueAsString(response);
String jsonResponse = toJson(response);
sendResponse(exchange, 200, "application/json", jsonResponse);
} catch (Exception e) {
log.error("Error generating /list response", e);
Expand All @@ -245,7 +242,7 @@ private void handleEval(HttpExchange exchange) throws IOException {
try {
InputStream requestBody = exchange.getRequestBody();
var requestBodyString = new String(requestBody.readAllBytes(), StandardCharsets.UTF_8);
EvalRequest request = JSON_MAPPER.readValue(requestBodyString, EvalRequest.class);
EvalRequest request = fromJson(requestBodyString, EvalRequest.class);

// Validate evaluator exists
RemoteEval eval = evals.get(request.getName());
Expand Down Expand Up @@ -543,22 +540,22 @@ private void setEvalSpanAttributes(
spanAttrs.put("generation", braintrustGeneration);
}
evalSpan.setAttribute(PARENT, braintrustParent.toParentValue())
.setAttribute("braintrust.span_attributes", json(spanAttrs))
.setAttribute("braintrust.input_json", json(Map.of("input", datasetCase.input())))
.setAttribute("braintrust.expected_json", json(datasetCase.expected()));
.setAttribute("braintrust.span_attributes", toJson(spanAttrs))
.setAttribute("braintrust.input_json", toJson(Map.of("input", datasetCase.input())))
.setAttribute("braintrust.expected_json", toJson(datasetCase.expected()));

if (datasetCase.origin().isPresent()) {
evalSpan.setAttribute("braintrust.origin", json(datasetCase.origin().get()));
evalSpan.setAttribute("braintrust.origin", toJson(datasetCase.origin().get()));
}
if (!datasetCase.tags().isEmpty()) {
evalSpan.setAttribute(
AttributeKey.stringArrayKey("braintrust.tags"), datasetCase.tags());
}
if (!datasetCase.metadata().isEmpty()) {
evalSpan.setAttribute("braintrust.metadata", json(datasetCase.metadata()));
evalSpan.setAttribute("braintrust.metadata", toJson(datasetCase.metadata()));
}
evalSpan.setAttribute(
"braintrust.output_json", json(Map.of("output", taskResult.result())));
"braintrust.output_json", toJson(Map.of("output", taskResult.result())));
}

private void setTaskSpanAttributes(
Expand All @@ -575,10 +572,10 @@ private void setTaskSpanAttributes(
}

taskSpan.setAttribute(PARENT, braintrustParent.toParentValue())
.setAttribute("braintrust.span_attributes", json(taskSpanAttrs))
.setAttribute("braintrust.input_json", json(Map.of("input", datasetCase.input())))
.setAttribute("braintrust.span_attributes", toJson(taskSpanAttrs))
.setAttribute("braintrust.input_json", toJson(Map.of("input", datasetCase.input())))
.setAttribute(
"braintrust.output_json", json(Map.of("output", taskResult.result())));
"braintrust.output_json", toJson(Map.of("output", taskResult.result())));
}

private void setScoreSpanAttributes(
Expand All @@ -594,10 +591,10 @@ private void setScoreSpanAttributes(
scoreSpanAttrs.put("generation", braintrustGeneration);
}

var scoresJson = json(scorerScores);
var scoresJson = toJson(scorerScores);
scoreSpan
.setAttribute(PARENT, braintrustParent.toParentValue())
.setAttribute("braintrust.span_attributes", json(scoreSpanAttrs))
.setAttribute("braintrust.span_attributes", toJson(scoreSpanAttrs))
.setAttribute("braintrust.output_json", scoresJson)
.setAttribute("braintrust.scores", scoresJson);
}
Expand Down Expand Up @@ -625,9 +622,9 @@ private void sendProgressEvent(
progressData.put("format", "code");
progressData.put("output_type", "completion");
progressData.put("event", "json_delta");
progressData.put("data", JSON_MAPPER.writeValueAsString(taskResult));
progressData.put("data", toJson(taskResult));

String progressJson = JSON_MAPPER.writeValueAsString(progressData);
String progressJson = toJson(progressData);
sendSSEEvent(os, "progress", progressJson);
}

Expand Down Expand Up @@ -663,21 +660,13 @@ private void sendSummaryEvent(
summary.put("scores", scoresWithMeta);
summary.put("metrics", Map.of());

sendSSEEvent(os, "summary", JSON_MAPPER.writeValueAsString(summary));
sendSSEEvent(os, "summary", toJson(summary));
}

private void sendDoneEvent(OutputStream os) throws IOException {
sendSSEEvent(os, "done", "");
}

private String json(Object o) {
try {
return JSON_MAPPER.writeValueAsString(o);
} catch (Exception e) {
throw new RuntimeException("Failed to serialize to JSON", e);
}
}

private void sendResponse(
HttpExchange exchange, int statusCode, String contentType, String body)
throws IOException {
Expand Down Expand Up @@ -934,8 +923,7 @@ private RequestContext getBraintrust(HttpExchange exchange, RequestContext conte
private void sendErrorResponse(HttpExchange exchange, int statusCode, String message)
throws IOException {
Map<String, String> error = Map.of("error", message);
String json = JSON_MAPPER.writeValueAsString(error);
sendResponse(exchange, statusCode, "application/json", json);
sendResponse(exchange, statusCode, "application/json", toJson(error));
}

/**
Expand Down
Loading