Skip to content

support custom serialization for datasets, remote evals, and scorers #136

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

Open
Andrew Kent (realark) wants to merge 3 commits into
base: main
Choose a base branch
from
+1,020 −30
Open

support custom serialization for datasets, remote evals, and scorers #136

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
51a3b7d
custom serialization for external datasets
realark Jul 7, 2026
8f7413b
custom serialization for remote evals
realark Jul 7, 2026
6abf086
custom serialization for remote scorers
realark Jul 7, 2026
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
180 changes: 180 additions & 0 deletions braintrust-sdk/src/main/java/dev/braintrust/Braintrust.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
import java.net.URI;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import lombok.Getter;
Expand Down Expand Up @@ -171,16 +172,150 @@ public <INPUT, OUTPUT> Eval.Builder<INPUT, OUTPUT> evalBuilder() {
Eval.builder().config(this.config).apiClient(this.openApiClient);
}

/**
* Fetch the latest version of a dataset from Braintrust, using the default project from
* configuration.
*
* @deprecated the {@code INPUT} and {@code OUTPUT} type parameters are not applied at runtime:
* case values are returned as raw JSON-decoded objects (e.g. {@code LinkedHashMap}), and
* accessing them as {@code INPUT}/{@code OUTPUT} will throw {@link ClassCastException}
* unless those types are {@code Object} or {@code Map}. Use {@link #fetchDataset(String,
* Class, Class)} instead.
*/
@Deprecated
public <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> fetchDataset(String datasetName) {
return fetchDataset(datasetName, null);
}

/**
* Fetch the latest version of a dataset from Braintrust, deserializing each case's {@code
* input} and {@code expected} values into the given types.
*
* <p>Values are deserialized using the SDK's shared Jackson mapper (see {@link
* dev.braintrust.json.BraintrustJsonMapper}), which may be customized via {@link
* dev.braintrust.json.BraintrustJsonMapper#configure}. For full control over per-value
* conversion (custom mappers, fixups, generic types), use {@link #fetchDataset(String,
* Function, Function)}.
*
* <p>The returned dataset preserves experiment-to-dataset linking when used with {@link Eval}:
* experiments created by {@code Eval.run()} are stamped with this dataset's id and version.
*
* @param datasetName the name of the dataset within the configured project
* @param inputClass the type to deserialize each case's {@code input} into
* @param outputClass the type to deserialize each case's {@code expected} into
* @return a typed view of the remote dataset; cases are fetched lazily, and deserialization
* errors surface during iteration
*/
public <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> fetchDataset(
String datasetName, Class<INPUT> inputClass, Class<OUTPUT> outputClass) {
return fetchDataset(datasetName, null, inputClass, outputClass);
}

/**
* Fetch a specific version of a dataset from Braintrust, deserializing each case's {@code
* input} and {@code expected} values into the given types.
*
* <p>See {@link #fetchDataset(String, Class, Class)} for deserialization and experiment-linking
* semantics.
*
* @param datasetName the name of the dataset within the configured project
* @param datasetVersion the dataset version to pin, or null to fetch the latest version upon
* every cursor open
* @param inputClass the type to deserialize each case's {@code input} into
* @param outputClass the type to deserialize each case's {@code expected} into
* @return a typed view of the remote dataset; cases are fetched lazily, and deserialization
* errors surface during iteration
*/
public <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> fetchDataset(
String datasetName,
@Nullable String datasetVersion,
Class<INPUT> inputClass,
Class<OUTPUT> outputClass) {
return Dataset.fetchFromBraintrust(
openApiClient,
resolveProjectName(),
datasetName,
datasetVersion,
inputClass,
outputClass);
}

/**
* Fetch a specific version of a dataset from Braintrust, using the default project from
* configuration.
*
* @deprecated the {@code INPUT} and {@code OUTPUT} type parameters are not applied at runtime;
* see {@link #fetchDataset(String)}. Use {@link #fetchDataset(String, String, Class,
* Class)} instead.
*/
@Deprecated
public <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> fetchDataset(
String datasetName, @Nullable String datasetVersion) {
return Dataset.fetchFromBraintrust(
openApiClient, resolveProjectName(), datasetName, datasetVersion);
}

/**
* Fetch the latest version of a dataset from Braintrust, converting each case's {@code input}
* and {@code expected} values with the supplied converter functions.
*
* <p>Converters receive the raw JSON-decoded value (typically a {@code Map<String, Object>},
* {@code List}, {@code String}, {@code Number}, {@code Boolean}, or null) and are responsible
* for producing the typed value. This gives the caller full control over deserialization — e.g.
* a custom Jackson {@code ObjectMapper}, generic types, or row fixups:
*
* <pre>{@code
* Dataset<MyInput, MyOutput> ds = braintrust.fetchDataset(
* "golden-cases",
* raw -> myMapper.convertValue(raw, MyInput.class),
* raw -> myMapper.convertValue(raw, MyOutput.class));
* }</pre>
*
* <p>The returned dataset preserves experiment-to-dataset linking when used with {@link Eval}:
* experiments created by {@code Eval.run()} are stamped with this dataset's id and version.
*
* @param datasetName the name of the dataset within the configured project
* @param inputConverter converts each case's raw {@code input} value; must tolerate null
* @param outputConverter converts each case's raw {@code expected} value; must tolerate null
* @return a typed view of the remote dataset; cases are fetched lazily, and converter errors
* surface during iteration
*/
public <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> fetchDataset(
String datasetName,
Function<Object, INPUT> inputConverter,
Function<Object, OUTPUT> outputConverter) {
return fetchDataset(datasetName, null, inputConverter, outputConverter);
}

/**
* Fetch a specific version of a dataset from Braintrust, converting each case's {@code input}
* and {@code expected} values with the supplied converter functions.
*
* <p>See {@link #fetchDataset(String, Function, Function)} for converter and experiment-linking
* semantics.
*
* @param datasetName the name of the dataset within the configured project
* @param datasetVersion the dataset version to pin, or null to fetch the latest version upon
* every cursor open
* @param inputConverter converts each case's raw {@code input} value; must tolerate null
* @param outputConverter converts each case's raw {@code expected} value; must tolerate null
* @return a typed view of the remote dataset; cases are fetched lazily, and converter errors
* surface during iteration
*/
public <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> fetchDataset(
String datasetName,
@Nullable String datasetVersion,
Function<Object, INPUT> inputConverter,
Function<Object, OUTPUT> outputConverter) {
return Dataset.fetchFromBraintrust(
openApiClient,
resolveProjectName(),
datasetName,
datasetVersion,
inputConverter,
outputConverter);
}

/**
* Fetch a scorer from Braintrust by slug, using the default project from configuration.
*
Expand All @@ -203,6 +338,30 @@ public <INPUT, OUTPUT> Scorer<INPUT, OUTPUT> fetchScorer(
return Scorer.fetchFromBraintrust(openApiClient, resolveProjectName(), scorerSlug, version);
}

/**
* Fetch a scorer from Braintrust by slug, using the default project from configuration, with a
* custom converter for serializing scorer argument values.
*
* <p>By default, argument values ({@code input}, {@code output}, {@code expected}, {@code
* metadata}, {@code parameters}) are serialized with the SDK's shared Jackson mapper (see
* {@link dev.braintrust.json.BraintrustJsonMapper}). Use this variant to control serialization,
* e.g. with a custom {@code ObjectMapper}: {@code fetchScorer(slug, null,
* myMapper::valueToTree)}.
*
* @param scorerSlug the unique slug identifier for the scorer
* @param version optional version of the scorer to fetch
* @param argumentConverter converts each argument value into a JSON-serializable form (e.g. a
* Jackson {@code JsonNode}, {@code Map}, or scalar); must tolerate null
* @return a Scorer that invokes the remote function
*/
public <INPUT, OUTPUT> Scorer<INPUT, OUTPUT> fetchScorer(
String scorerSlug,
@Nullable String version,
Function<Object, Object> argumentConverter) {
return Scorer.fetchFromBraintrust(
openApiClient, resolveProjectName(), scorerSlug, version, argumentConverter);
}

/**
* Fetch a scorer from Braintrust by project name and slug.
*
Expand All @@ -216,6 +375,27 @@ public <INPUT, OUTPUT> Scorer<INPUT, OUTPUT> fetchScorer(
return Scorer.fetchFromBraintrust(openApiClient, projectName, scorerSlug, version);
}

/**
* Fetch a scorer from Braintrust by project name and slug, with a custom converter for
* serializing scorer argument values. See {@link #fetchScorer(String, String, Function)} for
* converter semantics.
*
* @param projectName the name of the project containing the scorer
* @param scorerSlug the unique slug identifier for the scorer
* @param version optional version of the scorer to fetch
* @param argumentConverter converts each argument value into a JSON-serializable form; must
* tolerate null
* @return a Scorer that invokes the remote function
*/
public <INPUT, OUTPUT> Scorer<INPUT, OUTPUT> fetchScorer(
String projectName,
String scorerSlug,
@Nullable String version,
Function<Object, Object> argumentConverter) {
return Scorer.fetchFromBraintrust(
openApiClient, projectName, scorerSlug, version, argumentConverter);
}

/**
* Resolve the default project name from config. If only a project ID is configured, looks it up
* via the API. Mirrors the behavior of the old hand-rolled client.
Expand Down
52 changes: 37 additions & 15 deletions braintrust-sdk/src/main/java/dev/braintrust/devserver/Devserver.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,7 +347,7 @@ private void handleEval(HttpExchange exchange) throws IOException {
@SuppressWarnings({"unchecked", "rawtypes"})
private <I, O> void handleStreamingEval(
HttpExchange exchange,
RemoteEval eval,
RemoteEval<I, O> eval,
EvalRequest request,
RequestContext context,
List<Scorer<I, O>> remoteScorers)
Expand Down Expand Up @@ -418,11 +418,13 @@ private <I, O> void handleStreamingEval(

// NOTE: this code is serial but written in a thread-safe manner to support
// concurrent dataset fetching and eval execution
extractDataset(request, apiClient)
extractDataset(
request,
apiClient,
eval.getInputConverter(),
eval.getOutputConverter())
.forEach(
rawDataset -> {
final DatasetCase<I, O> datasetCase =
(DatasetCase<I, O>) rawDataset;
datasetCase -> {
var evalSpan =
tracer.spanBuilder("eval")
.setNoParent()
Expand Down Expand Up @@ -630,7 +632,12 @@ private <I, O> void handleExperimentSnapshot(
.config(braintrust.config())
.apiClient(apiClient)
.projectId(projectId)
.dataset((Dataset<I, O>) extractDataset(request, apiClient))
.dataset(
extractDataset(
request,
apiClient,
eval.getInputConverter(),
eval.getOutputConverter()))
.task(eval.getTask())
.scorers(allScorers.toArray(new Scorer[0]))
.parameters(eval.getParameters())
Expand Down Expand Up @@ -1283,31 +1290,41 @@ private static Optional<ParentInfo> extractPlaygroundParent(EvalRequest request)
* @throws IllegalStateException if no dataset specification is provided
* @throws IllegalArgumentException if dataset or project is not found
*/
private static Dataset<?, ?> extractDataset(
EvalRequest request, BraintrustOpenApiClient apiClient) {
private static <I, O> Dataset<I, O> extractDataset(
EvalRequest request,
BraintrustOpenApiClient apiClient,
java.util.function.Function<Object, I> inputConverter,
java.util.function.Function<Object, O> outputConverter) {
EvalRequest.DataSpec dataSpec = request.getData();

if (dataSpec.getData() != null && !dataSpec.getData().isEmpty()) {
// Method 1: Inline data
List<DatasetCase> cases = new ArrayList<>();
List<DatasetCase<I, O>> cases = new ArrayList<>();
for (EvalRequest.EvalCaseData caseData : dataSpec.getData()) {
DatasetCase datasetCase =
DatasetCase<I, O> datasetCase =
DatasetCase.of(
caseData.getInput(),
caseData.getExpected(),
inputConverter.apply(caseData.getInput()),
outputConverter.apply(caseData.getExpected()),
caseData.getTags() != null ? caseData.getTags() : List.of(),
caseData.getMetadata() != null ? caseData.getMetadata() : Map.of());
cases.add(datasetCase);
}
return Dataset.of(cases.toArray(new DatasetCase[0]));
@SuppressWarnings("unchecked")
DatasetCase<I, O>[] caseArray = cases.toArray(new DatasetCase[0]);
return Dataset.of(caseArray);
} else if (dataSpec.getProjectName() != null && dataSpec.getDatasetName() != null) {
// Method 2: Fetch by project name and dataset name
log.debug(
"Fetching dataset from Braintrust: project={}, dataset={}",
dataSpec.getProjectName(),
dataSpec.getDatasetName());
return Dataset.fetchFromBraintrust(
apiClient, dataSpec.getProjectName(), dataSpec.getDatasetName(), null);
apiClient,
dataSpec.getProjectName(),
dataSpec.getDatasetName(),
null,
inputConverter,
outputConverter);
} else if (dataSpec.getDatasetId() != null) {
// Method 3: Fetch by dataset ID
log.debug("Fetching dataset from Braintrust by ID: {}", dataSpec.getDatasetId());
Expand All @@ -1325,7 +1342,12 @@ private static Optional<ParentInfo> extractPlaygroundParent(EvalRequest request)
fetchedDatasetName);

return Dataset.fetchFromBraintrust(
apiClient, fetchedProjectName, fetchedDatasetName, null);
apiClient,
fetchedProjectName,
fetchedDatasetName,
null,
inputConverter,
outputConverter);
} else {
throw new IllegalStateException("No dataset specification provided");
}
Expand Down
Loading
Loading