From 23cc5ac9b744110a57d9b8c857a2fe5b31205cad Mon Sep 17 00:00:00 2001 From: Andrew Kent Date: Tue, 7 Jul 2026 00:16:29 -0600 Subject: [PATCH] fix custom serialization bugs in several places of the codebase we lost type info due to erasure which made it impossible to ser/de complex objects which fetching json from the braintrust api. Our examples used simple types jackson could handle out of the box, which is why we didn't hit this before. This change introduces type info into the public apis for fetching objects from braintrust. It also includes options for custom serializers (and by default uses the global braintrust mapper) affected: - remote datasets - remote scorers - devserver --- .../main/java/dev/braintrust/Braintrust.java | 193 ++++++++++++++++++ .../dev/braintrust/devserver/Devserver.java | 53 +++-- .../dev/braintrust/devserver/RemoteEval.java | 89 +++++++- .../java/dev/braintrust/eval/Dataset.java | 70 ++++++- .../eval/DatasetBrainstoreImpl.java | 85 +++++++- .../main/java/dev/braintrust/eval/Scorer.java | 55 ++++- .../braintrust/eval/ScorerBrainstoreImpl.java | 72 ++++++- .../braintrust/json/BraintrustJsonMapper.java | 12 ++ .../braintrust/devserver/DevserverTest.java | 106 ++++++++++ .../braintrust/devserver/RemoteEvalTest.java | 69 +++++++ .../eval/DatasetBrainstoreImplTest.java | 163 +++++++++++++++ ...ScorerBrainstoreImplSerializationTest.java | 153 ++++++++++++++ .../examples/RemoteEvalExample.java | 2 +- 13 files changed, 1082 insertions(+), 40 deletions(-) create mode 100644 braintrust-sdk/src/test/java/dev/braintrust/devserver/RemoteEvalTest.java create mode 100644 braintrust-sdk/src/test/java/dev/braintrust/eval/ScorerBrainstoreImplSerializationTest.java diff --git a/braintrust-sdk/src/main/java/dev/braintrust/Braintrust.java b/braintrust-sdk/src/main/java/dev/braintrust/Braintrust.java index 01bfda81..b9f95bf7 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/Braintrust.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/Braintrust.java @@ -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; @@ -171,16 +172,150 @@ public Eval.Builder 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 Dataset 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. + * + *

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)}. + * + *

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 Dataset fetchDataset( + String datasetName, Class inputClass, Class 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. + * + *

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 Dataset fetchDataset( + String datasetName, + @Nullable String datasetVersion, + Class inputClass, + Class 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 Dataset 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. + * + *

Converters receive the raw JSON-decoded value (typically a {@code Map}, + * {@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: + * + * 

{@code
+     * Dataset ds = braintrust.fetchDataset(
+     *         "golden-cases",
+     *         raw -> myMapper.convertValue(raw, MyInput.class),
+     *         raw -> myMapper.convertValue(raw, MyOutput.class));
+     * }
+ * + *

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 Dataset fetchDataset( + String datasetName, + Function inputConverter, + Function 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. + * + *

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 Dataset fetchDataset( + String datasetName, + @Nullable String datasetVersion, + Function inputConverter, + Function outputConverter) { + return Dataset.fetchFromBraintrust( + openApiClient, + resolveProjectName(), + datasetName, + datasetVersion, + inputConverter, + outputConverter); + } + /** * Fetch a scorer from Braintrust by slug, using the default project from configuration. * @@ -203,6 +338,39 @@ public Scorer fetchScorer( return Scorer.fetchFromBraintrust(openApiClient, resolveProjectName(), scorerSlug, version); } + /** + * Fetch a scorer from Braintrust by slug, using the default project from configuration, with + * custom converters that transform the {@code input}, {@code output}, and {@code expected} + * scorer argument values into a JSON-serializable form. + * + *

By default, argument values are serialized with the SDK's shared mapper ({@link + * dev.braintrust.json.BraintrustJsonMapper}). Use this variant to control serialization, e.g. + * with a custom {@code ObjectMapper}: {@code fetchScorer(slug, null, myMapper::valueToTree, + * myMapper::valueToTree)}. Case {@code metadata} and eval {@code parameters} are always + * serialized with the OpenAPI client's mapper. + * + * @param scorerSlug the unique slug identifier for the scorer + * @param version optional version of the scorer to fetch + * @param inputConverter converts each case's {@code input} value into a JSON-serializable form + * (e.g. a Jackson {@code JsonNode}, {@code Map}, or scalar); never invoked with null + * @param outputConverter converts the task {@code output} and case {@code expected} values into + * a JSON-serializable form; never invoked with null + * @return a Scorer that invokes the remote function + */ + public Scorer fetchScorer( + String scorerSlug, + @Nullable String version, + Function inputConverter, + Function outputConverter) { + return Scorer.fetchFromBraintrust( + openApiClient, + resolveProjectName(), + scorerSlug, + version, + inputConverter, + outputConverter); + } + /** * Fetch a scorer from Braintrust by project name and slug. * @@ -216,6 +384,31 @@ public Scorer fetchScorer( return Scorer.fetchFromBraintrust(openApiClient, projectName, scorerSlug, version); } + /** + * Fetch a scorer from Braintrust by project name and slug, with custom converters that + * transform the {@code input}, {@code output}, and {@code expected} scorer argument values into + * a JSON-serializable form. See {@link #fetchScorer(String, String, Function, 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 inputConverter converts each case's {@code input} value into a JSON-serializable form; + * never invoked with null + * @param outputConverter converts the task {@code output} and case {@code expected} values into + * a JSON-serializable form; never invoked with null + * @return a Scorer that invokes the remote function + */ + public Scorer fetchScorer( + String projectName, + String scorerSlug, + @Nullable String version, + Function inputConverter, + Function outputConverter) { + return Scorer.fetchFromBraintrust( + openApiClient, projectName, scorerSlug, version, inputConverter, outputConverter); + } + /** * 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. diff --git a/braintrust-sdk/src/main/java/dev/braintrust/devserver/Devserver.java b/braintrust-sdk/src/main/java/dev/braintrust/devserver/Devserver.java index 5a5ba9c9..1cb5a4ee 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/devserver/Devserver.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/devserver/Devserver.java @@ -30,6 +30,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; import java.util.function.Consumer; +import java.util.function.Function; import java.util.regex.Pattern; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -347,7 +348,7 @@ private void handleEval(HttpExchange exchange) throws IOException { @SuppressWarnings({"unchecked", "rawtypes"}) private void handleStreamingEval( HttpExchange exchange, - RemoteEval eval, + RemoteEval eval, EvalRequest request, RequestContext context, List> remoteScorers) @@ -418,11 +419,13 @@ private 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 datasetCase = - (DatasetCase) rawDataset; + datasetCase -> { var evalSpan = tracer.spanBuilder("eval") .setNoParent() @@ -630,7 +633,12 @@ private void handleExperimentSnapshot( .config(braintrust.config()) .apiClient(apiClient) .projectId(projectId) - .dataset((Dataset) extractDataset(request, apiClient)) + .dataset( + extractDataset( + request, + apiClient, + eval.getInputConverter(), + eval.getOutputConverter())) .task(eval.getTask()) .scorers(allScorers.toArray(new Scorer[0])) .parameters(eval.getParameters()) @@ -1283,23 +1291,28 @@ private static Optional 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 Dataset extractDataset( + EvalRequest request, + BraintrustOpenApiClient apiClient, + Function inputConverter, + Function outputConverter) { EvalRequest.DataSpec dataSpec = request.getData(); if (dataSpec.getData() != null && !dataSpec.getData().isEmpty()) { // Method 1: Inline data - List cases = new ArrayList<>(); + List> cases = new ArrayList<>(); for (EvalRequest.EvalCaseData caseData : dataSpec.getData()) { - DatasetCase datasetCase = + DatasetCase 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[] 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( @@ -1307,7 +1320,12 @@ private static Optional extractPlaygroundParent(EvalRequest request) 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()); @@ -1325,7 +1343,12 @@ private static Optional 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"); } diff --git a/braintrust-sdk/src/main/java/dev/braintrust/devserver/RemoteEval.java b/braintrust-sdk/src/main/java/dev/braintrust/devserver/RemoteEval.java index 9a2075c6..f1d272b0 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/devserver/RemoteEval.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/devserver/RemoteEval.java @@ -1,11 +1,10 @@ package dev.braintrust.devserver; -import dev.braintrust.eval.DatasetCase; import dev.braintrust.eval.ParameterDef; -import dev.braintrust.eval.Parameters; import dev.braintrust.eval.Scorer; import dev.braintrust.eval.Task; import dev.braintrust.eval.TaskResult; +import dev.braintrust.json.BraintrustJsonMapper; import java.util.*; import java.util.function.Function; import javax.annotation.Nonnull; @@ -42,6 +41,79 @@ public class RemoteEval { /** Optional parameter definitions that can be configured from the UI */ @Singular @Nonnull private final List> parameters; + /** + * Converts raw JSON-decoded {@code input} values from eval requests into {@code INPUT}. + * Defaults to an unchecked cast of the raw value. + */ + @lombok.Builder.Default @Nonnull + private final Function inputConverter = uncheckedCast(); + + /** + * Converts raw JSON-decoded {@code expected} values from eval requests into {@code OUTPUT}. + * Defaults to an unchecked cast of the raw value. + */ + @lombok.Builder.Default @Nonnull + private final Function outputConverter = uncheckedCast(); + + /** + * Create a builder for an eval whose case {@code input} and {@code expected} values are + * deserialized into the given types using the SDK's shared Jackson mapper (see {@link + * dev.braintrust.json.BraintrustJsonMapper}). + * + * 

{@code
+     * var eval = RemoteEval.builder(MyInput.class, MyOutput.class)
+     *         .name("my-eval")
+     *         .taskFunction(input -> ...)
+     *         .build();
+     * }
+ * + * @param inputClass the type to deserialize each case's {@code input} into + * @param outputClass the type to deserialize each case's {@code expected} into + */ + public static Builder builder( + Class inputClass, Class outputClass) { + return builder( + BraintrustJsonMapper.converter(inputClass), + BraintrustJsonMapper.converter(outputClass)); + } + + /** + * Create a builder for an eval whose case {@code input} and {@code expected} values are + * converted with the supplied functions. Converters receive the raw JSON-decoded value + * (typically a {@code Map}, {@code List}, {@code String}, {@code Number}, + * {@code Boolean}, or null) and must tolerate null. Most callers should prefer {@link + * #builder(Class, Class)}; use this variant for generic types or custom deserialization. + * + * @param inputConverter converts each case's raw {@code input} value + * @param outputConverter converts each case's raw {@code expected} value + */ + public static Builder builder( + Function inputConverter, Function outputConverter) { + return new Builder() + .inputConverter(Objects.requireNonNull(inputConverter)) + .outputConverter(Objects.requireNonNull(outputConverter)); + } + + /** + * Create a builder with no input/expected deserialization: case values are passed through as + * raw JSON-decoded objects (e.g. {@code LinkedHashMap}). + * + * @deprecated the {@code INPUT} and {@code OUTPUT} type parameters are not applied at runtime, + * and the task/scorers will receive raw JSON objects, throwing {@link ClassCastException} + * unless those types are {@code Object} or {@code Map}. Use {@link #builder(Class, Class)} + * or {@link #builder(Function, Function)} instead. + */ + @Deprecated + public static Builder builder() { + return RemoteEval.builder(uncheckedCast(), uncheckedCast()); + } + + /** Legacy behavior: pass the raw JSON-decoded value through via an unchecked cast. */ + @SuppressWarnings("unchecked") + private static Function uncheckedCast() { + return raw -> (T) raw; + } + public static class Builder { /** * Convenience builder method to create a RemoteEval with a simple task function. @@ -51,20 +123,17 @@ public static class Builder { */ public Builder taskFunction(Function taskFn) { return task( - new Task<>() { - @Override - public TaskResult apply( - DatasetCase datasetCase, Parameters parameters) - throws Exception { - var result = taskFn.apply(datasetCase.input()); - return new TaskResult<>(result, datasetCase, parameters); - } + (datasetCase, parameters) -> { + var result = taskFn.apply(datasetCase.input()); + return new TaskResult<>(result, datasetCase, parameters); }); } /** Build the RemoteEval */ public RemoteEval build() { var result = internalBuild(); + Objects.requireNonNull(result.getInputConverter(), "inputConverter"); + Objects.requireNonNull(result.getOutputConverter(), "outputConverter"); // Validate parameter names are unique var seen = new HashSet(); for (var param : result.getParameters()) { diff --git a/braintrust-sdk/src/main/java/dev/braintrust/eval/Dataset.java b/braintrust-sdk/src/main/java/dev/braintrust/eval/Dataset.java index ee42952b..671ff612 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/eval/Dataset.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/eval/Dataset.java @@ -6,6 +6,7 @@ import java.util.List; import java.util.Optional; import java.util.function.Consumer; +import java.util.function.Function; import javax.annotation.Nullable; import javax.annotation.concurrent.NotThreadSafe; @@ -74,11 +75,77 @@ static Dataset fetchFromBraintrust( apiClient.openApiClient(), projectName, datasetName, datasetVersion); } + /** + * Fetch a dataset from Braintrust. + * + * @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}). Use + * {@link #fetchFromBraintrust(BraintrustOpenApiClient, String, String, String, Class, + * Class)} instead. + */ + @Deprecated static Dataset fetchFromBraintrust( BraintrustOpenApiClient apiClient, String projectName, String datasetName, @Nullable String datasetVersion) { + var datasetId = resolveDatasetId(apiClient, projectName, datasetName); + return new DatasetBrainstoreImpl<>(apiClient, datasetId, datasetVersion); + } + + /** + * Fetch a dataset from Braintrust, deserializing each case's {@code input} and {@code expected} + * values into the given types using the SDK's shared Jackson mapper (see {@link + * dev.braintrust.json.BraintrustJsonMapper}). + * + * @param apiClient the Braintrust API client + * @param projectName the project containing the dataset + * @param datasetName the name of the dataset within the 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 + */ + static Dataset fetchFromBraintrust( + BraintrustOpenApiClient apiClient, + String projectName, + String datasetName, + @Nullable String datasetVersion, + Class inputClass, + Class outputClass) { + var datasetId = resolveDatasetId(apiClient, projectName, datasetName); + return new DatasetBrainstoreImpl<>( + apiClient, datasetId, datasetVersion, inputClass, outputClass); + } + + /** + * Fetch a dataset from Braintrust, converting each case's {@code input} and {@code expected} + * values with the supplied converter functions. Converters receive the raw JSON-decoded value + * (typically a {@code Map}, {@code List}, {@code String}, {@code Number}, + * {@code Boolean}, or null) and must tolerate null. + * + * @param apiClient the Braintrust API client + * @param projectName the project containing the dataset + * @param datasetName the name of the dataset within the 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 + * @param outputConverter converts each case's raw {@code expected} value + */ + static Dataset fetchFromBraintrust( + BraintrustOpenApiClient apiClient, + String projectName, + String datasetName, + @Nullable String datasetVersion, + Function inputConverter, + Function outputConverter) { + var datasetId = resolveDatasetId(apiClient, projectName, datasetName); + return new DatasetBrainstoreImpl<>( + apiClient, datasetId, datasetVersion, inputConverter, outputConverter); + } + + private static String resolveDatasetId( + BraintrustOpenApiClient apiClient, String projectName, String datasetName) { var datasetsApi = new DatasetsApi(apiClient); var objects = datasetsApi @@ -101,7 +168,6 @@ static Dataset fetchFromBraintrust( + " datasets"); } - var dataset = objects.get(0); - return new DatasetBrainstoreImpl<>(apiClient, dataset.getId().toString(), datasetVersion); + return objects.get(0).getId().toString(); } } diff --git a/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java b/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java index 60ecfc48..ad1dad88 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/eval/DatasetBrainstoreImpl.java @@ -2,9 +2,11 @@ import dev.braintrust.api.BraintrustApiClient; import dev.braintrust.api.BraintrustOpenApiClient; +import dev.braintrust.json.BraintrustJsonMapper; import dev.braintrust.openapi.api.DatasetsApi; import dev.braintrust.openapi.model.FetchEventsRequest; import java.util.*; +import java.util.function.Function; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -14,16 +16,72 @@ public class DatasetBrainstoreImpl implements Dataset inputConverter; + private final Function outputConverter; @Deprecated public DatasetBrainstoreImpl( BraintrustApiClient apiClient, String datasetId, @Nullable String datasetVersion) { - this(apiClient.openApiClient(), datasetId, datasetVersion, 512); + this(apiClient.openApiClient(), datasetId, datasetVersion); } + /** + * @deprecated the {@code INPUT} and {@code OUTPUT} type parameters are not applied at runtime: + * case values are returned as raw JSON-decoded objects. Use {@link + * #DatasetBrainstoreImpl(BraintrustOpenApiClient, String, String, Class, Class)} instead. + */ + @Deprecated public DatasetBrainstoreImpl( BraintrustOpenApiClient apiClient, String datasetId, @Nullable String datasetVersion) { - this(apiClient, datasetId, datasetVersion, 512); + this(apiClient, datasetId, datasetVersion, 512, uncheckedCast(), uncheckedCast()); + } + + /** + * Create a dataset whose case {@code input} and {@code expected} values are deserialized into + * the given types using the SDK's shared Jackson mapper (see {@link BraintrustJsonMapper}). + * + * @param apiClient the Braintrust API client + * @param datasetId the Braintrust dataset id + * @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 + */ + public DatasetBrainstoreImpl( + BraintrustOpenApiClient apiClient, + String datasetId, + @Nullable String datasetVersion, + Class inputClass, + Class outputClass) { + this( + apiClient, + datasetId, + datasetVersion, + 512, + BraintrustJsonMapper.converter(inputClass), + BraintrustJsonMapper.converter(outputClass)); + } + + /** + * Create a dataset whose case {@code input} and {@code expected} values are converted with the + * supplied functions. Converters receive the raw JSON-decoded value (typically a {@code + * Map}, {@code List}, {@code String}, {@code Number}, {@code Boolean}, or + * null). + * + * @param apiClient the Braintrust API client + * @param datasetId the Braintrust dataset id + * @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 + */ + public DatasetBrainstoreImpl( + BraintrustOpenApiClient apiClient, + String datasetId, + @Nullable String datasetVersion, + Function inputConverter, + Function outputConverter) { + this(apiClient, datasetId, datasetVersion, 512, inputConverter, outputConverter); } DatasetBrainstoreImpl( @@ -31,10 +89,28 @@ public DatasetBrainstoreImpl( String datasetId, @Nullable String datasetVersion, int batchSize) { + this(apiClient, datasetId, datasetVersion, batchSize, uncheckedCast(), uncheckedCast()); + } + + DatasetBrainstoreImpl( + BraintrustOpenApiClient apiClient, + String datasetId, + @Nullable String datasetVersion, + int batchSize, + Function inputConverter, + Function outputConverter) { this.apiClient = apiClient; this.datasetId = datasetId; this.batchSize = batchSize; this.pinnedVersion = datasetVersion; + this.inputConverter = Objects.requireNonNull(inputConverter); + this.outputConverter = Objects.requireNonNull(outputConverter); + } + + /** Legacy behavior: pass the raw JSON-decoded value through via an unchecked cast. */ + @SuppressWarnings("unchecked") + private static Function uncheckedCast() { + return raw -> (T) raw; } @Override @@ -98,7 +174,6 @@ private class BrainstoreCursor implements Cursor> { } @Override - @SuppressWarnings("unchecked") public Optional> next() { if (closed) { throw new IllegalStateException("Cursor is closed"); @@ -114,8 +189,8 @@ public Optional> next() { var event = currentBatch.get(currentIndex++); - INPUT input = (INPUT) event.getInput(); - OUTPUT expected = (OUTPUT) event.getExpected(); + INPUT input = inputConverter.apply(event.getInput()); + OUTPUT expected = outputConverter.apply(event.getExpected()); var metadataObj = event.getMetadata(); // InsertProjectLogsEventMetadata extends HashMap. Jackson stores diff --git a/braintrust-sdk/src/main/java/dev/braintrust/eval/Scorer.java b/braintrust-sdk/src/main/java/dev/braintrust/eval/Scorer.java index 1b62fa25..300a5607 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/eval/Scorer.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/eval/Scorer.java @@ -111,6 +111,59 @@ static Scorer fetchFromBraintrust( String projectName, String scorerSlug, @Nullable String version) { + return new ScorerBrainstoreImpl<>( + apiClient, resolveFunctionId(apiClient, projectName, scorerSlug, version), version); + } + + /** + * Fetch a scorer from Braintrust by project name and slug, with custom converters that + * transform the {@code input}, {@code output}, and {@code expected} scorer argument values into + * a JSON-serializable form. + * + *

By default ({@link #fetchFromBraintrust(BraintrustOpenApiClient, String, String, + * String)}), argument values are serialized with the SDK's shared mapper ({@link + * dev.braintrust.json.BraintrustJsonMapper}). Use this variant to control serialization, e.g. + * with a custom {@code ObjectMapper}: + * + * 

{@code
+     * Scorer scorer = Scorer.fetchFromBraintrust(
+     *         client, project, slug, null, myMapper::valueToTree, myMapper::valueToTree);
+     * }
+ * + *

Case {@code metadata} and eval {@code parameters} are always serialized with the OpenAPI + * client's mapper. + * + * @param apiClient the API client to use + * @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 inputConverter converts each case's {@code input} value into a JSON-serializable form + * (e.g. a Jackson {@code JsonNode}, {@code Map}, or scalar); never invoked with null + * @param outputConverter converts the task {@code output} and case {@code expected} values into + * a JSON-serializable form; never invoked with null + * @return a Scorer that invokes the remote function + * @throws RuntimeException if the scorer is not found + */ + static Scorer fetchFromBraintrust( + BraintrustOpenApiClient apiClient, + String projectName, + String scorerSlug, + @Nullable String version, + Function inputConverter, + Function outputConverter) { + return new ScorerBrainstoreImpl<>( + apiClient, + resolveFunctionId(apiClient, projectName, scorerSlug, version), + version, + inputConverter, + outputConverter); + } + + private static String resolveFunctionId( + BraintrustOpenApiClient apiClient, + String projectName, + String scorerSlug, + @Nullable String version) { var functionsApi = new FunctionsApi(apiClient); var objects = functionsApi @@ -133,6 +186,6 @@ static Scorer fetchFromBraintrust( "Scorer not found: project=" + projectName + ", slug=" + scorerSlug); } - return new ScorerBrainstoreImpl<>(apiClient, objects.get(0).getId().toString(), version); + return objects.get(0).getId().toString(); } } diff --git a/braintrust-sdk/src/main/java/dev/braintrust/eval/ScorerBrainstoreImpl.java b/braintrust-sdk/src/main/java/dev/braintrust/eval/ScorerBrainstoreImpl.java index ec491262..e5b4ffa3 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/eval/ScorerBrainstoreImpl.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/eval/ScorerBrainstoreImpl.java @@ -4,6 +4,7 @@ import dev.braintrust.BraintrustUtils; import dev.braintrust.api.BraintrustApiClient; import dev.braintrust.api.BraintrustOpenApiClient; +import dev.braintrust.json.BraintrustJsonMapper; import dev.braintrust.openapi.JSON; import dev.braintrust.openapi.api.FunctionsApi; import dev.braintrust.openapi.api.ProjectsApi; @@ -39,9 +40,26 @@ public class ScorerBrainstoreImpl implements Scorer { private static final ObjectMapper MAPPER = new JSON().getMapper(); + /** + * Default conversion for user-supplied input/output/expected scorer argument values: serializes + * to a JSON tree using the SDK's shared mapper ({@link + * dev.braintrust.json.BraintrustJsonMapper}), matching how typed dataset values are + * deserialized and how eval span data is logged. To apply custom serialization rules, supply + * custom input/output converters. + */ + private static final java.util.function.Function DEFAULT_CONVERTER = + value -> BraintrustJsonMapper.get().valueToTree(value); + + @SuppressWarnings("unchecked") + private static java.util.function.Function defaultConverter() { + return (java.util.function.Function) DEFAULT_CONVERTER; + } + private final BraintrustOpenApiClient apiClient; private final String functionId; private final @Nullable String version; + private final java.util.function.Function inputConverter; + private final java.util.function.Function outputConverter; private final AtomicReference braintrustFunction = new AtomicReference<>(null); private final AtomicReference scorerName = new AtomicReference<>(null); @@ -52,7 +70,10 @@ public ScorerBrainstoreImpl( } /** - * Create a new remote scorer. + * Create a new remote scorer. Input, output, and expected values are serialized with the SDK's + * shared mapper ({@link dev.braintrust.json.BraintrustJsonMapper}); use {@link + * #ScorerBrainstoreImpl(BraintrustOpenApiClient, String, String, java.util.function.Function, + * java.util.function.Function)} to customize serialization. * * @param apiClient the API client to use for invoking the function * @param functionId braintrust function id @@ -61,9 +82,38 @@ public ScorerBrainstoreImpl( */ public ScorerBrainstoreImpl( BraintrustOpenApiClient apiClient, String functionId, @Nullable String version) { + this(apiClient, functionId, version, defaultConverter(), defaultConverter()); + } + + /** + * Create a new remote scorer with custom input/output converters. + * + * @param apiClient the API client to use for invoking the function + * @param functionId braintrust function id + * @param version optional version of the function to invoke. null always invokes latest + * version. + * @param inputConverter converts each case's {@code input} value into a JSON-serializable form + * (e.g. a Jackson {@code JsonNode}, {@code Map}, or scalar) before the invoke request is + * sent; never invoked with null + * @param outputConverter converts the task {@code output} and case {@code expected} values into + * a JSON-serializable form; never invoked with null + */ + public ScorerBrainstoreImpl( + BraintrustOpenApiClient apiClient, + String functionId, + @Nullable String version, + java.util.function.Function inputConverter, + java.util.function.Function outputConverter) { this.apiClient = apiClient; this.functionId = functionId; this.version = version; + this.inputConverter = java.util.Objects.requireNonNull(inputConverter); + this.outputConverter = java.util.Objects.requireNonNull(outputConverter); + } + + private static Object convert( + java.util.function.Function converter, @Nullable T value) { + return value == null ? null : converter.apply(value); } @Override @@ -74,13 +124,23 @@ public String getName() { @Override public List score(TaskResult taskResult) { + // Convert user-supplied input/output/expected values into a JSON-serializable form. + // The default converters serialize with the SDK's shared BraintrustJsonMapper, matching + // typed dataset deserialization and eval span logging; custom converters can be supplied + // to control serialization. Metadata and parameters are plain JSON maps and are always + // serialized with the OpenAPI client's mapper. The OpenAPI client's own mapper then + // writes the converted values verbatim, keeping wire-protocol handling of the request + // envelope unchanged. var scorerArgs = new java.util.LinkedHashMap(); - scorerArgs.put("input", taskResult.datasetCase().input()); - scorerArgs.put("output", taskResult.result()); - scorerArgs.put("expected", taskResult.datasetCase().expected()); - scorerArgs.put("metadata", taskResult.datasetCase().metadata()); + scorerArgs.put("input", convert(inputConverter, taskResult.datasetCase().input())); + scorerArgs.put("output", convert(outputConverter, taskResult.result())); + scorerArgs.put("expected", convert(outputConverter, taskResult.datasetCase().expected())); + scorerArgs.put( + "metadata", convert(MAPPER::valueToTree, taskResult.datasetCase().metadata())); if (!taskResult.parameters().isEmpty()) { - scorerArgs.put("parameters", taskResult.parameters().getMerged()); + scorerArgs.put( + "parameters", + convert(MAPPER::valueToTree, taskResult.parameters().getMerged())); } var invoke = new InvokeApi().input(scorerArgs).version(version); diff --git a/braintrust-sdk/src/main/java/dev/braintrust/json/BraintrustJsonMapper.java b/braintrust-sdk/src/main/java/dev/braintrust/json/BraintrustJsonMapper.java index 91952df2..3e872463 100644 --- a/braintrust-sdk/src/main/java/dev/braintrust/json/BraintrustJsonMapper.java +++ b/braintrust-sdk/src/main/java/dev/braintrust/json/BraintrustJsonMapper.java @@ -8,7 +8,9 @@ import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; import java.util.ArrayList; import java.util.List; +import java.util.Objects; import java.util.function.Consumer; +import java.util.function.Function; import lombok.SneakyThrows; /** Centralized ObjectMapper for the Braintrust SDK. */ @@ -69,6 +71,16 @@ public static T fromJson(String jsonString, Class targetClass) { return get().readValue(jsonString, targetClass); } + /** + * Returns a converter function that deserializes raw JSON-decoded values (e.g. {@code + * Map}, {@code List}, {@code String}, {@code Number}, {@code Boolean}) into the + * given type using the shared mapper. Null values pass through as null. + */ + public static Function converter(Class targetClass) { + Objects.requireNonNull(targetClass); + return raw -> raw == null ? null : get().convertValue(raw, targetClass); + } + static synchronized void reset() { instance = null; configurers.clear(); diff --git a/braintrust-sdk/src/test/java/dev/braintrust/devserver/DevserverTest.java b/braintrust-sdk/src/test/java/dev/braintrust/devserver/DevserverTest.java index d76dc2e5..173fdb69 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/devserver/DevserverTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/devserver/DevserverTest.java @@ -39,6 +39,12 @@ class DevserverTest { private static final String REMOTE_EVAL_NAME = "food-type-classifier"; private static final String PARAM_EVAL_NAME = "param-eval"; + private static final String TYPED_EVAL_NAME = "typed-eval"; + + public record TypedInput(String prompt, int n) {} + + public record TypedOutput(String answer) {} + private static final String TASK_ERROR_EVAL_NAME = "task-error-eval"; private static final String SCORER_ERROR_EVAL_NAME = "scorer-error-eval"; private static final BraintrustUtils.Parent PLAYGROUND_PARENT = @@ -154,8 +160,23 @@ public dev.braintrust.eval.TaskResult apply( .scorer(Scorer.of("static_scorer", (expected, result) -> 1.0)) .build(); + // Typed eval: input/expected are deserialized into records via builder(Class, Class) + RemoteEval typedEval = + RemoteEval.builder(TypedInput.class, TypedOutput.class) + .name(TYPED_EVAL_NAME) + .taskFunction(input -> new TypedOutput(input.prompt() + ":" + input.n())) + .scorer( + Scorer.of( + "typed_scorer", + (expected, result) -> + expected.answer().equals(result.answer()) + ? 1.0 + : 0.0)) + .build(); + server = Devserver.builder() + .registerEval(typedEval) .config(testHarness.braintrust().config()) .registerEval(testEval) .registerEval(paramEval) @@ -209,6 +230,91 @@ void testHealthCheck() throws Exception { assertEquals("text/plain", response.headers().firstValue("Content-Type").orElse("")); } + @Test + void testStreamingEvalWithTypedInputOutput() throws Exception { + // Inline case data arrives as raw JSON; the typed eval (built via + // RemoteEval.builder(Class, Class)) must deserialize it into records before + // invoking the task and scorers. + EvalRequest evalRequest = new EvalRequest(); + evalRequest.setName(TYPED_EVAL_NAME); + evalRequest.setStream(true); + + EvalRequest.DataSpec dataSpec = new EvalRequest.DataSpec(); + EvalRequest.EvalCaseData case1 = new EvalRequest.EvalCaseData(); + case1.setInput(Map.of("prompt", "apple", "n", 2)); + case1.setExpected(Map.of("answer", "apple:2")); + dataSpec.setData(List.of(case1)); + evalRequest.setData(dataSpec); + + Map parentSpec = + Map.of( + "object_type", PLAYGROUND_PARENT.type(), + "object_id", PLAYGROUND_PARENT.id(), + "propagated_event", + Map.of("span_attributes", Map.of("generation", "test-gen-typed"))); + evalRequest.setParent(parentSpec); + + String requestBody = JSON_MAPPER.writeValueAsString(evalRequest); + + HttpURLConnection conn = + (HttpURLConnection) new URI(TEST_URL + "/eval").toURL().openConnection(); + conn.setRequestMethod("POST"); + conn.setRequestProperty("Content-Type", "application/json"); + conn.setRequestProperty("x-bt-auth-token", testHarness.braintrustApiKey()); + conn.setRequestProperty("x-bt-project-id", TestHarness.defaultProjectId()); + conn.setRequestProperty("x-bt-org-name", TestHarness.defaultOrgName()); + conn.setDoOutput(true); + conn.getOutputStream().write(requestBody.getBytes(StandardCharsets.UTF_8)); + conn.getOutputStream().flush(); + + assertEquals(200, conn.getResponseCode()); + + List> events = readSseEvents(conn); + + List> progressEvents = + events.stream().filter(e -> "progress".equals(e.get("event"))).toList(); + List> summaryEvents = + events.stream().filter(e -> "summary".equals(e.get("event"))).toList(); + assertEquals(1, progressEvents.size(), "Should have 1 progress event"); + assertEquals(1, summaryEvents.size(), "Should have 1 summary event"); + + // The task only compiles/runs against typed records: output proves input was + // deserialized into TypedInput rather than passed through as a LinkedHashMap. + JsonNode progressData = JSON_MAPPER.readTree(progressEvents.get(0).get("data")); + JsonNode taskResult = JSON_MAPPER.readTree(progressData.get("data").asText()); + assertEquals("apple:2", taskResult.get("answer").asText()); + + // The scorer compares typed expected vs typed result: 1.0 proves expected was + // deserialized into TypedOutput. + JsonNode summaryData = JSON_MAPPER.readTree(summaryEvents.get(0).get("data")); + JsonNode typedScorer = summaryData.get("scores").get("typed_scorer"); + assertEquals(1.0, typedScorer.get("score").asDouble(), 0.001); + } + + private static List> readSseEvents(HttpURLConnection conn) + throws Exception { + try (BufferedReader reader = + new BufferedReader( + new InputStreamReader(conn.getInputStream(), StandardCharsets.UTF_8))) { + List> events = new ArrayList<>(); + String line; + String currentEvent = null; + StringBuilder currentData = new StringBuilder(); + while ((line = reader.readLine()) != null) { + if (line.startsWith("event: ")) { + currentEvent = line.substring(7); + } else if (line.startsWith("data: ")) { + currentData.append(line.substring(6)); + } else if (line.isEmpty() && currentEvent != null) { + events.add(Map.of("event", currentEvent, "data", currentData.toString())); + currentEvent = null; + currentData = new StringBuilder(); + } + } + return events; + } + } + @Test void testStreamingEval() throws Exception { // Create eval request with inline data using EvalRequest types diff --git a/braintrust-sdk/src/test/java/dev/braintrust/devserver/RemoteEvalTest.java b/braintrust-sdk/src/test/java/dev/braintrust/devserver/RemoteEvalTest.java new file mode 100644 index 00000000..787d8bc4 --- /dev/null +++ b/braintrust-sdk/src/test/java/dev/braintrust/devserver/RemoteEvalTest.java @@ -0,0 +1,69 @@ +package dev.braintrust.devserver; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.LinkedHashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; + +public class RemoteEvalTest { + + public record MyInput(String prompt, int temperature) {} + + public record MyOutput(String answer) {} + + @Test + void builderWithClassesConvertsRawJsonValues() { + RemoteEval eval = + RemoteEval.builder(MyInput.class, MyOutput.class) + .name("typed-eval") + .taskFunction(input -> new MyOutput(input.prompt())) + .build(); + + Map rawInput = new LinkedHashMap<>(); + rawInput.put("prompt", "hello"); + rawInput.put("temperature", 7); + + MyInput input = eval.getInputConverter().apply(rawInput); + assertEquals(new MyInput("hello", 7), input); + + MyOutput expected = eval.getOutputConverter().apply(Map.of("answer", "world")); + assertEquals(new MyOutput("world"), expected); + + // null values pass through as null + assertNull(eval.getInputConverter().apply(null)); + assertNull(eval.getOutputConverter().apply(null)); + } + + @Test + void builderWithConvertersUsesSuppliedFunctions() { + RemoteEval eval = + RemoteEval.builder( + raw -> { + @SuppressWarnings("unchecked") + var map = (Map) raw; + return new MyInput( + (String) map.get("prompt"), + ((Number) map.get("temperature")).intValue()); + }, + raw -> String.valueOf(raw)) + .name("converter-eval") + .taskFunction(input -> input.prompt()) + .build(); + + MyInput input = eval.getInputConverter().apply(Map.of("prompt", "hi", "temperature", 3)); + assertEquals(new MyInput("hi", 3), input); + assertEquals("42", eval.getOutputConverter().apply(42)); + } + + @Test + @SuppressWarnings("deprecation") + void deprecatedBuilderPassesRawValuesThrough() { + RemoteEval eval = + RemoteEval.builder().name("legacy-eval").taskFunction(input -> input).build(); + + Map raw = Map.of("prompt", "hello"); + assertSame(raw, eval.getInputConverter().apply(raw)); + assertSame(raw, eval.getOutputConverter().apply(raw)); + } +} diff --git a/braintrust-sdk/src/test/java/dev/braintrust/eval/DatasetBrainstoreImplTest.java b/braintrust-sdk/src/test/java/dev/braintrust/eval/DatasetBrainstoreImplTest.java index 17d712a1..e8b91fb4 100644 --- a/braintrust-sdk/src/test/java/dev/braintrust/eval/DatasetBrainstoreImplTest.java +++ b/braintrust-sdk/src/test/java/dev/braintrust/eval/DatasetBrainstoreImplTest.java @@ -306,6 +306,169 @@ void testTagsPopulatedFromDatasetRow() { assertEquals("unit-test", tags.get(0)); } + public record MyInput(String prompt, int temperature) {} + + public record MyOutput(String answer) {} + + private void stubTypedRow() { + wireMock.stubFor( + post(urlEqualTo("/v1/dataset/" + datasetId + "/fetch")) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "events": [ + { + "object_type": "dataset", + "dataset_id": "%s", + "id": "typed-row-1", + "_xact_id": "1", + "created": "2024-01-01T00:00:00Z", + "input": {"prompt": "hello", "temperature": 7}, + "expected": {"answer": "world"} + } + ], + "cursor": null + } + """ + .formatted(datasetId)))); + } + + @Test + void testTypedDeserializationWithClasses() { + stubTypedRow(); + + Dataset dataset = + new DatasetBrainstoreImpl<>( + apiClient, datasetId, "test-version", MyInput.class, MyOutput.class); + + List> cases = new ArrayList<>(); + dataset.forEach(cases::add); + + assertEquals(1, cases.size()); + assertEquals(MyInput.class, cases.get(0).input().getClass()); + assertEquals(new MyInput("hello", 7), cases.get(0).input()); + assertEquals(new MyOutput("world"), cases.get(0).expected()); + } + + @Test + void testTypedDeserializationWithConverters() { + stubTypedRow(); + + Dataset dataset = + new DatasetBrainstoreImpl<>( + apiClient, + datasetId, + "test-version", + raw -> { + @SuppressWarnings("unchecked") + var map = (Map) raw; + return new MyInput( + (String) map.get("prompt"), + ((Number) map.get("temperature")).intValue()); + }, + raw -> { + @SuppressWarnings("unchecked") + var map = (Map) raw; + return (String) map.get("answer"); + }); + + List> cases = new ArrayList<>(); + dataset.forEach(cases::add); + + assertEquals(1, cases.size()); + assertEquals(new MyInput("hello", 7), cases.get(0).input()); + assertEquals("world", cases.get(0).expected()); + } + + @Test + void testTypedDeserializationNullValues() { + wireMock.stubFor( + post(urlEqualTo("/v1/dataset/" + datasetId + "/fetch")) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "events": [ + { + "object_type": "dataset", + "dataset_id": "%s", + "id": "null-row-1", + "_xact_id": "1", + "created": "2024-01-01T00:00:00Z", + "input": {"prompt": "hello", "temperature": 7} + } + ], + "cursor": null + } + """ + .formatted(datasetId)))); + + Dataset dataset = + new DatasetBrainstoreImpl<>( + apiClient, datasetId, "test-version", MyInput.class, MyOutput.class); + + List> cases = new ArrayList<>(); + dataset.forEach(cases::add); + + assertEquals(1, cases.size()); + assertEquals(new MyInput("hello", 7), cases.get(0).input()); + assertNull(cases.get(0).expected()); + } + + @Test + void testTypedFetchFromBraintrust() { + String projectName = "test-project"; + String datasetName = "typed-dataset"; + + wireMock.stubFor( + get(urlPathEqualTo("/v1/dataset")) + .withQueryParam("project_name", equalTo(projectName)) + .withQueryParam("dataset_name", equalTo(datasetName)) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "objects": [ + { + "id": "%s", + "project_id": "00000000-0000-0000-0000-000000000456", + "name": "typed-dataset", + "_xact_id": "1", + "created": "2024-01-01T00:00:00Z" + } + ] + } + """ + .formatted(datasetId)))); + stubTypedRow(); + + Dataset dataset = + Dataset.fetchFromBraintrust( + apiClient, projectName, datasetName, "1", MyInput.class, MyOutput.class); + + // typed fetch must still return the Brainstore-backed impl so Eval.run() links + // experiment <-> dataset + assertInstanceOf(DatasetBrainstoreImpl.class, dataset); + assertEquals(datasetId, dataset.id()); + + List> cases = new ArrayList<>(); + dataset.forEach(cases::add); + + assertEquals(1, cases.size()); + assertEquals(new MyInput("hello", 7), cases.get(0).input()); + assertEquals(new MyOutput("world"), cases.get(0).expected()); + } + @Test void testFetchFromBraintrustNotFound() { String projectName = "test-project"; diff --git a/braintrust-sdk/src/test/java/dev/braintrust/eval/ScorerBrainstoreImplSerializationTest.java b/braintrust-sdk/src/test/java/dev/braintrust/eval/ScorerBrainstoreImplSerializationTest.java new file mode 100644 index 00000000..d75f5f4d --- /dev/null +++ b/braintrust-sdk/src/test/java/dev/braintrust/eval/ScorerBrainstoreImplSerializationTest.java @@ -0,0 +1,153 @@ +package dev.braintrust.eval; + +import static com.github.tomakehurst.wiremock.client.WireMock.*; +import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig; +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.github.tomakehurst.wiremock.junit5.WireMockExtension; +import dev.braintrust.api.BraintrustOpenApiClient; +import dev.braintrust.config.BraintrustConfig; +import dev.braintrust.json.BraintrustJsonMapper; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** + * Verifies that user-supplied task/dataset objects in scorer invoke payloads are serialized with + * the SDK's shared mapper ({@link BraintrustJsonMapper}) by default, and that custom input/output + * converters can be supplied to control serialization. + */ +public class ScorerBrainstoreImplSerializationTest { + + @RegisterExtension + static WireMockExtension wireMock = + WireMockExtension.newInstance().options(wireMockConfig().dynamicPort()).build(); + + private static final String FUNCTION_ID = "00000000-0000-0000-0000-000000000abc"; + private static final String PROJECT_ID = "00000000-0000-0000-0000-000000000def"; + + private BraintrustOpenApiClient apiClient; + + public record MyInput(String userPrompt, @JsonIgnore String secretKey) {} + + public record MyOutput(String finalAnswer) {} + + @BeforeEach + void beforeEach() { + wireMock.resetAll(); + var config = + BraintrustConfig.builder() + .apiKey("test-api-key") + .apiUrl("http://localhost:" + wireMock.getPort()) + .build(); + apiClient = BraintrustOpenApiClient.of(config); + + wireMock.stubFor( + get(urlPathEqualTo("/v1/function/" + FUNCTION_ID)) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "id": "%s", + "project_id": "%s", + "name": "my-scorer", + "slug": "my-scorer", + "_xact_id": "1", + "created": "2024-01-01T00:00:00Z", + "log_id": "p", + "org_id": "00000000-0000-0000-0000-000000000001", + "function_data": {"type": "code"} + } + """ + .formatted(FUNCTION_ID, PROJECT_ID)))); + + wireMock.stubFor( + get(urlPathEqualTo("/v1/project/" + PROJECT_ID)) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody( + """ + { + "id": "%s", + "org_id": "00000000-0000-0000-0000-000000000001", + "name": "test-project" + } + """ + .formatted(PROJECT_ID)))); + + wireMock.stubFor( + post(urlPathEqualTo("/v1/function/" + FUNCTION_ID + "/invoke")) + .willReturn( + aResponse() + .withStatus(200) + .withHeader("Content-Type", "application/json") + .withBody("0.5"))); + } + + @Test + void scorerArgsAreSerializedWithBraintrustJsonMapperByDefault() { + Scorer scorer = new ScorerBrainstoreImpl<>(apiClient, FUNCTION_ID, null); + + var datasetCase = + DatasetCase.of(new MyInput("what is 2+2?", "hunter2"), new MyOutput("four")); + var taskResult = new TaskResult<>(new MyOutput("4"), datasetCase); + + var scores = scorer.score(taskResult); + assertEquals(1, scores.size()); + assertEquals(0.5, scores.get(0).value(), 0.001); + + wireMock.verify( + postRequestedFor(urlPathEqualTo("/v1/function/" + FUNCTION_ID + "/invoke")) + // default (BraintrustJsonMapper) uses snake_case field names, + // matching typed dataset deserialization and eval span logging + .withRequestBody( + matchingJsonPath( + "$.input.input.user_prompt", equalTo("what is 2+2?"))) + .withRequestBody( + matchingJsonPath("$.input.output.final_answer", equalTo("4"))) + .withRequestBody( + matchingJsonPath("$.input.expected.final_answer", equalTo("four"))) + // @JsonIgnore fields must not leak into the payload + .withRequestBody(notMatching(".*secret.*"))); + } + + @Test + void customConvertersControlSerialization() { + // Converters that wrap the input/output values, proving user control over serialization. + Scorer scorer = + new ScorerBrainstoreImpl<>( + apiClient, + FUNCTION_ID, + null, + (MyInput input) -> + java.util.Map.of( + "wrapped", BraintrustJsonMapper.get().valueToTree(input)), + (MyOutput output) -> + java.util.Map.of( + "wrapped", BraintrustJsonMapper.get().valueToTree(output))); + + var datasetCase = + DatasetCase.of(new MyInput("what is 2+2?", "hunter2"), new MyOutput("four")); + var taskResult = new TaskResult<>(new MyOutput("4"), datasetCase); + + var scores = scorer.score(taskResult); + assertEquals(1, scores.size()); + assertEquals(0.5, scores.get(0).value(), 0.001); + + wireMock.verify( + postRequestedFor(urlPathEqualTo("/v1/function/" + FUNCTION_ID + "/invoke")) + .withRequestBody( + matchingJsonPath( + "$.input.input.wrapped.user_prompt", + equalTo("what is 2+2?"))) + .withRequestBody( + matchingJsonPath( + "$.input.output.wrapped.final_answer", equalTo("4")))); + } +} diff --git a/examples/remote-eval/src/main/java/dev/braintrust/examples/RemoteEvalExample.java b/examples/remote-eval/src/main/java/dev/braintrust/examples/RemoteEvalExample.java index ad83ae3f..36f57135 100644 --- a/examples/remote-eval/src/main/java/dev/braintrust/examples/RemoteEvalExample.java +++ b/examples/remote-eval/src/main/java/dev/braintrust/examples/RemoteEvalExample.java @@ -18,7 +18,7 @@ public static void main(String[] args) throws Exception { var openAIClient = BraintrustOpenAI.wrapOpenAI(openTelemetry, OpenAIOkHttpClient.fromEnv()); RemoteEval foodTypeEval = - RemoteEval.builder() + RemoteEval.builder(String.class, String.class) .name("food-type-classifier") .taskFunction( food -> {