diff --git a/eval/eval.go b/eval/eval.go index bd5fae6..d898437 100644 --- a/eval/eval.go +++ b/eval/eval.go @@ -89,6 +89,7 @@ type Opts[I, R any] struct { Metadata Metadata // Metadata to attach to the experiment Update bool // If true, append to existing experiment (default: false) Parallelism int // Number of goroutines (default: 1) + TrialCount int // Number of times to run each case (default: 1) Quiet bool // Suppress result output (default: false) } @@ -109,6 +110,10 @@ type Case[I, R any] struct { // Optional. Metadata map[string]interface{} + // TrialCount overrides Opts.TrialCount for this case. + // Optional. Values less than 1 use the evaluation default. + TrialCount int + // These fields are only set if the Case is part of a Dataset. // They link the eval result back to the source dataset row. ID string // Dataset record ID @@ -240,13 +245,15 @@ type eval[I, R any] struct { tracer oteltrace.Tracer startSpanOpt oteltrace.SpanStartOption goroutines int + trialCount int quiet bool } // nextCase is a wrapper for sending cases through a channel. type nextCase[I, R any] struct { - c Case[I, R] - iterErr error + c Case[I, R] + trialIndex int + iterErr error } // newEval creates a new eval executor from concrete parameters (low-level constructor). @@ -263,6 +270,7 @@ func newEval[I, R any]( scorers []Scorer[I, R], classifiers []Classifier[I, R], parallelism int, + trialCount int, quiet bool, ) *eval[I, R] { // Build parent span option @@ -278,6 +286,10 @@ func newEval[I, R any]( goroutines = 1 } + if trialCount < 1 { + trialCount = 1 + } + return &eval[I, R]{ session: s, parent: parent, @@ -293,6 +305,7 @@ func newEval[I, R any]( tracer: tracer, startSpanOpt: startSpanOpt, goroutines: goroutines, + trialCount: trialCount, quiet: quiet, } } @@ -331,6 +344,7 @@ func newEvalOpts[I, R any](ctx context.Context, s *auth.Session, tp *trace.Trace opts.Scorers, opts.Classifiers, opts.Parallelism, + opts.TrialCount, opts.Quiet, ), nil } @@ -373,7 +387,17 @@ func (e *eval[I, R]) run(ctx context.Context) (*Result, error) { close(nextCases) break } - nextCases <- nextCase[I, R]{c: c, iterErr: err} + if err != nil { + nextCases <- nextCase[I, R]{iterErr: err} + continue + } + trialCount := c.TrialCount + if trialCount < 1 { + trialCount = e.trialCount + } + for trialIndex := range trialCount { + nextCases <- nextCase[I, R]{c: c, trialIndex: trialIndex, iterErr: err} + } } // 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 } // otherwise let's run the case (using the existing span) - return e.runCase(ctx, span, nextCase.c) + return e.runCase(ctx, span, nextCase.c, nextCase.trialIndex) } // runCase orchestrates task + scorers for one case. -func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I, R]) error { +func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I, R], trialIndex int) error { // Set all non-output attributes upfront so they're captured even if the task fails. attrs := map[string]any{ "braintrust.span_attributes": evalSpanAttrs, @@ -449,7 +473,7 @@ func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I, span.SetAttributes(attribute.StringSlice("braintrust.tags", c.Tags)) } - taskResult, err := e.runTask(ctx, span, c) + taskResult, err := e.runTask(ctx, span, c, trialIndex) if err != nil { span.SetStatus(codes.Error, err.Error()) return err @@ -492,7 +516,7 @@ func (e *eval[I, R]) runCase(ctx context.Context, span oteltrace.Span, c Case[I, // runTask executes the task function and creates a task span. // Returns a TaskResult containing all task execution data. -func (e *eval[I, R]) runTask(ctx context.Context, evalSpan oteltrace.Span, c Case[I, R]) (TaskResult[I, R], error) { +func (e *eval[I, R]) runTask(ctx context.Context, evalSpan oteltrace.Span, c Case[I, R], trialIndex int) (TaskResult[I, R], error) { ctx, taskSpan := e.tracer.Start(ctx, "task", e.startSpanOpt) defer taskSpan.End() @@ -511,11 +535,12 @@ func (e *eval[I, R]) runTask(ctx context.Context, evalSpan oteltrace.Span, c Cas // Construct TaskHooks with both spans and case data hooks := &TaskHooks{ - Expected: c.Expected, - Metadata: c.Metadata, - Tags: c.Tags, - TaskSpan: taskSpan, - EvalSpan: evalSpan, + Expected: c.Expected, + Metadata: c.Metadata, + Tags: c.Tags, + TrialIndex: trialIndex, + TaskSpan: taskSpan, + EvalSpan: evalSpan, } // Call task with new signature @@ -942,6 +967,7 @@ func testNewEval[I, R any]( scorers, classifiers, parallelism, + 1, // trialCount=1 for tests true, // quiet=true for tests ) } diff --git a/eval/eval_test.go b/eval/eval_test.go index f27d27d..e0f88c8 100644 --- a/eval/eval_test.go +++ b/eval/eval_test.go @@ -4,6 +4,8 @@ import ( "context" "errors" "io" + "sync" + "sync/atomic" "testing" "github.com/stretchr/testify/assert" @@ -312,6 +314,60 @@ func TestNewEval_Parallelism(t *testing.T) { assert.Equal(t, 4, ute.eval.goroutines) } +func TestEval_TrialCount(t *testing.T) { + t.Parallel() + + cases := NewDataset([]Case[testInput, testOutput]{ + {Input: testInput{Value: "test1"}}, + {Input: testInput{Value: "test2"}}, + }) + + var calls atomic.Int64 + var mu sync.Mutex + trialIndices := map[string][]int{} + task := func(ctx context.Context, input testInput, hooks *TaskHooks) (TaskOutput[testOutput], error) { + calls.Add(1) + mu.Lock() + trialIndices[input.Value] = append(trialIndices[input.Value], hooks.TrialIndex) + mu.Unlock() + return TaskOutput[testOutput]{Value: testOutput{Result: input.Value}}, nil + } + + ute := newUnitTestEval(t, cases, task, nil, 1) + ute.eval.trialCount = 3 + + _, err := ute.eval.run(context.Background()) + require.NoError(t, err) + assert.Equal(t, int64(6), calls.Load()) + assert.ElementsMatch(t, []int{0, 1, 2}, trialIndices["test1"]) + assert.ElementsMatch(t, []int{0, 1, 2}, trialIndices["test2"]) + + spans := ute.exporter.Flush() + assert.Len(t, spans, 12) // 2 cases * 3 trials * (task + eval) +} + +func TestEval_CaseTrialCountOverridesDefault(t *testing.T) { + t.Parallel() + + cases := NewDataset([]Case[testInput, testOutput]{ + {Input: testInput{Value: "test1"}}, + {Input: testInput{Value: "test2"}, TrialCount: 4}, + }) + + var calls atomic.Int64 + task := T(func(ctx context.Context, input testInput) (testOutput, error) { + calls.Add(1) + return testOutput{Result: input.Value}, nil + }) + + ute := newUnitTestEval(t, cases, task, nil, 1) + ute.eval.trialCount = 2 + + _, err := ute.eval.run(context.Background()) + require.NoError(t, err) + assert.Equal(t, int64(6), calls.Load()) +} + func TestNewEval_DefaultParallelism(t *testing.T) { t.Parallel() @@ -872,6 +928,7 @@ func TestEval_ParallelWithIteratorErrors(t *testing.T) { }) ute := newUnitTestEval(t, generator, task, nil, 2) // parallel=2 + ute.eval.trialCount = 3 ctx := context.Background() result, err := ute.eval.run(ctx) @@ -882,8 +939,8 @@ func TestEval_ParallelWithIteratorErrors(t *testing.T) { assert.NotNil(t, result) spans := ute.exporter.Flush() - // No scorers: 3 successful cases * 2 spans (task+eval) + 1 iterator error span = 7 spans - assert.Len(t, spans, 7) + // No scorers: 3 successful cases * 3 trials * 2 spans (task+eval) + 1 iterator error span = 19 spans + assert.Len(t, spans, 19) } // customCases allows custom Next() implementation for testing diff --git a/eval/task.go b/eval/task.go index 77e96e2..6f57677 100644 --- a/eval/task.go +++ b/eval/task.go @@ -19,9 +19,10 @@ type TaskHooks struct { // Readonly fields. These aren't necessarily recommended to be included in the task function, // but are available for advanced use cases. - Expected any // Not usually used in tasks, so this is untyped - Metadata Metadata // Case metadata - Tags []string // Case tags + Expected any // Not usually used in tasks, so this is untyped + Metadata Metadata // Case metadata + Tags []string // Case tags + TrialIndex int // The index of the current trial (0-based) } // TaskOutput wraps the output value from a task. diff --git a/examples/internal/eval-trial-count/main.go b/examples/internal/eval-trial-count/main.go new file mode 100644 index 0000000..6d194aa --- /dev/null +++ b/examples/internal/eval-trial-count/main.go @@ -0,0 +1,61 @@ +// This example demonstrates running each eval case multiple times with TrialCount. +// This is useful for nondeterministic tasks where you want repeated measurements per input. +package main + +import ( + "context" + "fmt" + "log" + "strings" + "time" + + "go.opentelemetry.io/otel/sdk/trace" + + "github.com/braintrustdata/braintrust-sdk-go" + "github.com/braintrustdata/braintrust-sdk-go/eval" +) + +func main() { + tp := trace.NewTracerProvider() + defer tp.Shutdown(context.Background()) //nolint:errcheck + + bt, err := braintrust.New(tp, + braintrust.WithProject("go-sdk-examples"), + ) + if err != nil { + log.Fatalf("Error creating Braintrust client: %v", err) + } + + var taskRuns int + task := func(ctx context.Context, input string) (string, error) { + taskRuns++ + return strings.ToUpper(input), nil + } + + scorer := eval.NewScorer("exact_match", func(ctx context.Context, result eval.TaskResult[string, string]) (eval.Scores, error) { + if result.Output == result.Expected { + return eval.S(1), nil + } + return eval.S(0), nil + }) + + cases := []eval.Case[string, string]{ + {Input: "hello", Expected: "HELLO"}, // Uses global TrialCount (2) + {Input: "world", Expected: "WORLD", TrialCount: 3}, // Overrides global TrialCount + } + + evaluator := braintrust.NewEvaluator[string, string](bt) + result, err := evaluator.Run(context.Background(), eval.Opts[string, string]{ + Experiment: fmt.Sprintf("trial-count-demo-%d", time.Now().Unix()), + Dataset: eval.NewDataset(cases), + Task: eval.T(task), + Scorers: []eval.Scorer[string, string]{scorer}, + TrialCount: 2, + }) + if err != nil { + log.Fatalf("Eval failed: %v", err) + } + + permalink, _ := result.Permalink() + log.Printf("Eval complete after %d task runs: %s", taskRuns, permalink) +}