Skip to content

Commit f567198

Browse files
authored
feat(eval): add trial count support (#162)
Run eval cases multiple times with a global TrialCount option and allow each case to override the default. Expose the zero-based TrialIndex through task hooks so tasks can distinguish repeated executions. Public API additions: ```go // eval.Opts TrialCount int // eval.Case TrialCount int // eval.TaskHooks TrialIndex int ``` Example usage: ```go result, err := evaluator.Run(ctx, eval.Opts[string, string]{ Experiment: "trial-count-demo", Dataset: eval.NewDataset([]eval.Case[string, string]{ {Input: "hello", Expected: "HELLO"}, {Input: "world", Expected: "WORLD", TrialCount: 3}, // per-case override }), Task: eval.T(task), Scorers: []eval.Scorer[string, string]{scorer}, TrialCount: 2, // default trials per case }) ``` ```go task := func(ctx context.Context, input string, hooks *eval.TaskHooks) (eval.TaskOutput[string], error) { log.Printf("running trial %d", hooks.TrialIndex) return eval.TaskOutput[string]{Value: strings.ToUpper(input)}, nil } ```
1 parent 80e311d commit f567198

4 files changed

Lines changed: 162 additions & 17 deletions

File tree

eval/eval.go

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ type Opts[I, R any] struct {
8989
Metadata Metadata // Metadata to attach to the experiment
9090
Update bool // If true, append to existing experiment (default: false)
9191
Parallelism int // Number of goroutines (default: 1)
92+
TrialCount int // Number of times to run each case (default: 1)
9293
Quiet bool // Suppress result output (default: false)
9394
}
9495

@@ -109,6 +110,10 @@ type Case[I, R any] struct {
109110
// Optional.
110111
Metadata map[string]interface{}
111112

113+
// TrialCount overrides Opts.TrialCount for this case.
114+
// Optional. Values less than 1 use the evaluation default.
115+
TrialCount int
116+
112117
// These fields are only set if the Case is part of a Dataset.
113118
// They link the eval result back to the source dataset row.
114119
ID string // Dataset record ID
@@ -240,13 +245,15 @@ type eval[I, R any] struct {
240245
tracer oteltrace.Tracer
241246
startSpanOpt oteltrace.SpanStartOption
242247
goroutines int
248+
trialCount int
243249
quiet bool
244250
}
245251

246252
// nextCase is a wrapper for sending cases through a channel.
247253
type nextCase[I, R any] struct {
248-
c Case[I, R]
249-
iterErr error
254+
c Case[I, R]
255+
trialIndex int
256+
iterErr error
250257
}
251258

252259
// newEval creates a new eval executor from concrete parameters (low-level constructor).
@@ -263,6 +270,7 @@ func newEval[I, R any](
263270
scorers []Scorer[I, R],
264271
classifiers []Classifier[I, R],
265272
parallelism int,
273+
trialCount int,
266274
quiet bool,
267275
) *eval[I, R] {
268276
// Build parent span option
@@ -278,6 +286,10 @@ func newEval[I, R any](
278286
goroutines = 1
279287
}
280288

289+
if trialCount < 1 {
290+
trialCount = 1
291+
}
292+
281293
return &eval[I, R]{
282294
session: s,
283295
parent: parent,
@@ -293,6 +305,7 @@ func newEval[I, R any](
293305
tracer: tracer,
294306
startSpanOpt: startSpanOpt,
295307
goroutines: goroutines,
308+
trialCount: trialCount,
296309
quiet: quiet,
297310
}
298311
}
@@ -331,6 +344,7 @@ func newEvalOpts[I, R any](ctx context.Context, s *auth.Session, tp *trace.Trace
331344
opts.Scorers,
332345
opts.Classifiers,
333346
opts.Parallelism,
347+
opts.TrialCount,
334348
opts.Quiet,
335349
), nil
336350
}
@@ -373,7 +387,17 @@ func (e *eval[I, R]) run(ctx context.Context) (*Result, error) {
373387
close(nextCases)
374388
break
375389
}
376-
nextCases <- nextCase[I, R]{c: c, iterErr: err}
390+
if err != nil {
391+
nextCases <- nextCase[I, R]{iterErr: err}
392+
continue
393+
}
394+
trialCount := c.TrialCount
395+
if trialCount < 1 {
396+
trialCount = e.trialCount
397+
}
398+
for trialIndex := range trialCount {
399+
nextCases <- nextCase[I, R]{c: c, trialIndex: trialIndex, iterErr: err}
400+
}
377401
}
378402

379403
// Wait for all the goroutines to finish.
@@ -419,11 +443,11 @@ func (e *eval[I, R]) runNextCase(ctx context.Context, nextCase nextCase[I, R]) e
419443
}
420444

421445
// otherwise let's run the case (using the existing span)
422-
return e.runCase(ctx, span, nextCase.c)
446+
return e.runCase(ctx, span, nextCase.c, nextCase.trialIndex)
423447
}
424448

425449
// runCase orchestrates task + scorers for one case.
426-
func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I, R]) error {
450+
func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I, R], trialIndex int) error {
427451
// Set all non-output attributes upfront so they're captured even if the task fails.
428452
attrs := map[string]any{
429453
"braintrust.span_attributes": evalSpanAttrs,
@@ -449,7 +473,7 @@ func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I,
449473
span.SetAttributes(attribute.StringSlice("braintrust.tags", c.Tags))
450474
}
451475

452-
taskResult, err := e.runTask(ctx, span, c)
476+
taskResult, err := e.runTask(ctx, span, c, trialIndex)
453477
if err != nil {
454478
span.SetStatus(codes.Error, err.Error())
455479
return err
@@ -492,7 +516,7 @@ func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I,
492516

493517
// runTask executes the task function and creates a task span.
494518
// Returns a TaskResult containing all task execution data.
495-
func (e *eval[I, R]) runTask(ctx context.Context, evalSpan oteltrace.Span, c Case[I, R]) (TaskResult[I, R], error) {
519+
func (e *eval[I, R]) runTask(ctx context.Context, evalSpan oteltrace.Span, c Case[I, R], trialIndex int) (TaskResult[I, R], error) {
496520
ctx, taskSpan := e.tracer.Start(ctx, "task", e.startSpanOpt)
497521
defer taskSpan.End()
498522

@@ -511,11 +535,12 @@ func (e *eval[I, R]) runTask(ctx context.Context, evalSpan oteltrace.Span, c Cas
511535

512536
// Construct TaskHooks with both spans and case data
513537
hooks := &TaskHooks{
514-
Expected: c.Expected,
515-
Metadata: c.Metadata,
516-
Tags: c.Tags,
517-
TaskSpan: taskSpan,
518-
EvalSpan: evalSpan,
538+
Expected: c.Expected,
539+
Metadata: c.Metadata,
540+
Tags: c.Tags,
541+
TrialIndex: trialIndex,
542+
TaskSpan: taskSpan,
543+
EvalSpan: evalSpan,
519544
}
520545

521546
// Call task with new signature
@@ -942,6 +967,7 @@ func testNewEval[I, R any](
942967
scorers,
943968
classifiers,
944969
parallelism,
970+
1, // trialCount=1 for tests
945971
true, // quiet=true for tests
946972
)
947973
}

eval/eval_test.go

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import (
44
"context"
55
"errors"
66
"io"
7+
"sync"
8+
"sync/atomic"
79
"testing"
810

911
"github.com/stretchr/testify/assert"
@@ -312,6 +314,60 @@ func TestNewEval_Parallelism(t *testing.T) {
312314
assert.Equal(t, 4, ute.eval.goroutines)
313315
}
314316

317+
func TestEval_TrialCount(t *testing.T) {
318+
t.Parallel()
319+
320+
cases := NewDataset([]Case[testInput, testOutput]{
321+
{Input: testInput{Value: "test1"}},
322+
{Input: testInput{Value: "test2"}},
323+
})
324+
325+
var calls atomic.Int64
326+
var mu sync.Mutex
327+
trialIndices := map[string][]int{}
328+
task := func(ctx context.Context, input testInput, hooks *TaskHooks) (TaskOutput[testOutput], error) {
329+
calls.Add(1)
330+
mu.Lock()
331+
trialIndices[input.Value] = append(trialIndices[input.Value], hooks.TrialIndex)
332+
mu.Unlock()
333+
return TaskOutput[testOutput]{Value: testOutput{Result: input.Value}}, nil
334+
}
335+
336+
ute := newUnitTestEval(t, cases, task, nil, 1)
337+
ute.eval.trialCount = 3
338+
339+
_, err := ute.eval.run(context.Background())
340+
require.NoError(t, err)
341+
assert.Equal(t, int64(6), calls.Load())
342+
assert.ElementsMatch(t, []int{0, 1, 2}, trialIndices["test1"])
343+
assert.ElementsMatch(t, []int{0, 1, 2}, trialIndices["test2"])
344+
345+
spans := ute.exporter.Flush()
346+
assert.Len(t, spans, 12) // 2 cases * 3 trials * (task + eval)
347+
}
348+
349+
func TestEval_CaseTrialCountOverridesDefault(t *testing.T) {
350+
t.Parallel()
351+
352+
cases := NewDataset([]Case[testInput, testOutput]{
353+
{Input: testInput{Value: "test1"}},
354+
{Input: testInput{Value: "test2"}, TrialCount: 4},
355+
})
356+
357+
var calls atomic.Int64
358+
task := T(func(ctx context.Context, input testInput) (testOutput, error) {
359+
calls.Add(1)
360+
return testOutput{Result: input.Value}, nil
361+
})
362+
363+
ute := newUnitTestEval(t, cases, task, nil, 1)
364+
ute.eval.trialCount = 2
365+
366+
_, err := ute.eval.run(context.Background())
367+
require.NoError(t, err)
368+
assert.Equal(t, int64(6), calls.Load())
369+
}
370+
315371
func TestNewEval_DefaultParallelism(t *testing.T) {
316372
t.Parallel()
317373

@@ -872,6 +928,7 @@ func TestEval_ParallelWithIteratorErrors(t *testing.T) {
872928
})
873929

874930
ute := newUnitTestEval(t, generator, task, nil, 2) // parallel=2
931+
ute.eval.trialCount = 3
875932

876933
ctx := context.Background()
877934
result, err := ute.eval.run(ctx)
@@ -882,8 +939,8 @@ func TestEval_ParallelWithIteratorErrors(t *testing.T) {
882939
assert.NotNil(t, result)
883940

884941
spans := ute.exporter.Flush()
885-
// No scorers: 3 successful cases * 2 spans (task+eval) + 1 iterator error span = 7 spans
886-
assert.Len(t, spans, 7)
942+
// No scorers: 3 successful cases * 3 trials * 2 spans (task+eval) + 1 iterator error span = 19 spans
943+
assert.Len(t, spans, 19)
887944
}
888945

889946
// customCases allows custom Next() implementation for testing

eval/task.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,10 @@ type TaskHooks struct {
1919

2020
// Readonly fields. These aren't necessarily recommended to be included in the task function,
2121
// but are available for advanced use cases.
22-
Expected any // Not usually used in tasks, so this is untyped
23-
Metadata Metadata // Case metadata
24-
Tags []string // Case tags
22+
Expected any // Not usually used in tasks, so this is untyped
23+
Metadata Metadata // Case metadata
24+
Tags []string // Case tags
25+
TrialIndex int // The index of the current trial (0-based)
2526
}
2627

2728
// TaskOutput wraps the output value from a task.
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
// This example demonstrates running each eval case multiple times with TrialCount.
2+
// This is useful for nondeterministic tasks where you want repeated measurements per input.
3+
package main
4+
5+
import (
6+
"context"
7+
"fmt"
8+
"log"
9+
"strings"
10+
"time"
11+
12+
"go.opentelemetry.io/otel/sdk/trace"
13+
14+
"github.com/braintrustdata/braintrust-sdk-go"
15+
"github.com/braintrustdata/braintrust-sdk-go/eval"
16+
)
17+
18+
func main() {
19+
tp := trace.NewTracerProvider()
20+
defer tp.Shutdown(context.Background()) //nolint:errcheck
21+
22+
bt, err := braintrust.New(tp,
23+
braintrust.WithProject("go-sdk-examples"),
24+
)
25+
if err != nil {
26+
log.Fatalf("Error creating Braintrust client: %v", err)
27+
}
28+
29+
var taskRuns int
30+
task := func(ctx context.Context, input string) (string, error) {
31+
taskRuns++
32+
return strings.ToUpper(input), nil
33+
}
34+
35+
scorer := eval.NewScorer("exact_match", func(ctx context.Context, result eval.TaskResult[string, string]) (eval.Scores, error) {
36+
if result.Output == result.Expected {
37+
return eval.S(1), nil
38+
}
39+
return eval.S(0), nil
40+
})
41+
42+
cases := []eval.Case[string, string]{
43+
{Input: "hello", Expected: "HELLO"}, // Uses global TrialCount (2)
44+
{Input: "world", Expected: "WORLD", TrialCount: 3}, // Overrides global TrialCount
45+
}
46+
47+
evaluator := braintrust.NewEvaluator[string, string](bt)
48+
result, err := evaluator.Run(context.Background(), eval.Opts[string, string]{
49+
Experiment: fmt.Sprintf("trial-count-demo-%d", time.Now().Unix()),
50+
Dataset: eval.NewDataset(cases),
51+
Task: eval.T(task),
52+
Scorers: []eval.Scorer[string, string]{scorer},
53+
TrialCount: 2,
54+
})
55+
if err != nil {
56+
log.Fatalf("Eval failed: %v", err)
57+
}
58+
59+
permalink, _ := result.Permalink()
60+
log.Printf("Eval complete after %d task runs: %s", taskRuns, permalink)
61+
}

0 commit comments

Comments
 (0)