Merge pull request #11 from braintrustdata/ark/eval-task-output

Rework evals api

Author: realark

README.md

@@ -52,9 +52,9 @@ Function<String, String> getFoodType =
 var eval = braintrust.<String, String>evalBuilder()
                 .name("java-eval-x-" + System.currentTimeMillis())
                 .cases(
-                        EvalCase.of("asparagus", "vegetable"),
-                        EvalCase.of("banana", "fruit"))
-                .task(getFoodType)
+                        DatasetCase.of("asparagus", "vegetable"),
+                        DatasetCase.of("banana", "fruit"))
+                .taskFunction(getFoodType)
                 .scorers(
                         Scorer.of(
                                 "fruit_scorer",

examples/src/main/java/dev/braintrust/examples/ExperimentExample.java

@@ -4,7 +4,7 @@
 import com.openai.models.ChatModel;
 import com.openai.models.chat.completions.ChatCompletionCreateParams;
 import dev.braintrust.Braintrust;
-import dev.braintrust.eval.EvalCase;
+import dev.braintrust.eval.DatasetCase;
 import dev.braintrust.eval.Scorer;
 import dev.braintrust.instrumentation.openai.BraintrustOpenAI;
 import java.util.function.Function;
@@ -37,10 +37,10 @@ public static void main(String[] args) throws Exception {
                         // will append new cases to
                         // the same experiment
                         .cases(
-                                EvalCase.of("strawberry", "fruit"),
-                                EvalCase.of("asparagus", "vegetable"),
-                                EvalCase.of("apple", "fruit"),
-                                EvalCase.of("banana", "fruit"))
+                                DatasetCase.of("strawberry", "fruit"),
+                                DatasetCase.of("asparagus", "vegetable"),
+                                DatasetCase.of("apple", "fruit"),
+                                DatasetCase.of("banana", "fruit"))
                         .taskFunction(getFoodType)
                         .scorers(
                                 Scorer.of(

src/main/java/dev/braintrust/eval/Dataset.java

@@ -0,0 +1,42 @@
+package dev.braintrust.eval;
+
+import java.util.List;
+import java.util.Optional;
+import javax.annotation.concurrent.NotThreadSafe;
+
+/**
+ * Datasets define the cases for evals. This interface provides a means of iterating through all
+ * cases of a particular dataset.
+ *
+ * <p>The most common implementations are in-memory datasets, and datasets fetched from the
+ * Braintrust API.
+ */
+public interface Dataset<INPUT, OUTPUT> {
+    Cursor<DatasetCase<INPUT, OUTPUT>> openCursor();
+
+    String id();
+
+    String version();
+
+    @NotThreadSafe
+    interface Cursor<CASE> extends AutoCloseable {
+        /**
+         * Fetch the next case. Returns empty if there are no more cases to fetch.
+         *
+         * <p>Implementations may make external requests to fetch data.
+         *
+         * <p>If this method is invoked after {@link #close()} an IllegalStateException will be
+         * thrown
+         */
+        Optional<CASE> next();
+
+        /** close all cursor resources */
+        void close();
+    }
+
+    /** Create an in-memory Dataset containing the provided cases. */
+    @SafeVarargs
+    static <INPUT, OUTPUT> Dataset<INPUT, OUTPUT> of(DatasetCase<INPUT, OUTPUT>... cases) {
+        return new DatasetInMemoryImpl<>(List.of(cases));
+    }
+}

src/main/java/dev/braintrust/eval/DatasetCase.java

@@ -0,0 +1,25 @@
+package dev.braintrust.eval;
+
+import java.util.List;
+import java.util.Map;
+import javax.annotation.Nonnull;
+
+/** A single row in a dataset. */
+public record DatasetCase<INPUT, OUTPUT>(
+        INPUT input,
+        OUTPUT expected,
+        @Nonnull List<String> tags,
+        @Nonnull Map<String, Object> metadata) {
+    public DatasetCase {
+        if (!metadata.isEmpty()) {
+            throw new RuntimeException("TODO: metadata support not yet implemented");
+        }
+        if (!tags.isEmpty()) {
+            throw new RuntimeException("TODO: tags support not yet implemented");
+        }
+    }
+
+    public static <INPUT, OUTPUT> DatasetCase<INPUT, OUTPUT> of(INPUT input, OUTPUT expected) {
+        return new DatasetCase<>(input, expected, List.of(), Map.of());
+    }
+}

src/main/java/dev/braintrust/eval/DatasetInMemoryImpl.java

@@ -0,0 +1,49 @@
+package dev.braintrust.eval;
+
+import java.util.List;
+import java.util.Optional;
+
+/** A dataset held entirely in memory */
+class DatasetInMemoryImpl<INPUT, OUTPUT> implements Dataset<INPUT, OUTPUT> {
+    private final List<DatasetCase<INPUT, OUTPUT>> cases;
+    private final String id;
+
+    DatasetInMemoryImpl(List<DatasetCase<INPUT, OUTPUT>> cases) {
+        this.cases = List.copyOf(cases);
+        id = "in-memory-dataset<" + this.cases.hashCode() + ">";
+    }
+
+    @Override
+    public String id() {
+        return id;
+    }
+
+    @Override
+    public String version() {
+        return "0";
+    }
+
+    @Override
+    public Cursor<DatasetCase<INPUT, OUTPUT>> openCursor() {
+        return new Cursor<>() {
+            int nextIndex = 0;
+            boolean closed = false;
+
+            @Override
+            public Optional<DatasetCase<INPUT, OUTPUT>> next() {
+                if (closed) {
+                    throw new IllegalStateException("this method may not be invoked after close");
+                } else if (nextIndex < cases.size()) {
+                    return Optional.of(cases.get(nextIndex++));
+                } else {
+                    return Optional.empty();
+                }
+            }
+
+            @Override
+            public void close() {
+                closed = true;
+            }
+        };
+    }
+}