diff --git a/ldai/judge/judge.go b/ldai/judge/judge.go index 608dd67a..1cad7d5d 100644 --- a/ldai/judge/judge.go +++ b/ldai/judge/judge.go @@ -21,7 +21,7 @@ type Config interface { type Tracker interface { TrackJudgeResponse(response datamodel.JudgeResponse) error - TrackUsage(usage ldai.TokenUsage) error + TrackTokens(usage ldai.TokenUsage) error } type StructuredResponse struct { @@ -91,7 +91,7 @@ func (j *Judge) Evaluate(input, output string, samplingRate float64) (*datamodel } if response.Usage.Total > 0 || response.Usage.Input > 0 || response.Usage.Output > 0 { - _ = j.tracker.TrackUsage(response.Usage) + _ = j.tracker.TrackTokens(response.Usage) } result := j.parseResponse(response.Content) diff --git a/ldai/judge/judge_test.go b/ldai/judge/judge_test.go index df0fd451..da0d7e55 100644 --- a/ldai/judge/judge_test.go +++ b/ldai/judge/judge_test.go @@ -51,7 +51,7 @@ func (m *mockTracker) TrackJudgeResponse(response datamodel.JudgeResponse) error return nil } -func (m *mockTracker) TrackUsage(usage ldai.TokenUsage) error { +func (m *mockTracker) TrackTokens(usage ldai.TokenUsage) error { m.usages = append(m.usages, usage) return nil } diff --git a/ldai/tracker.go b/ldai/tracker.go index 90655242..ff380ccb 100644 --- a/ldai/tracker.go +++ b/ldai/tracker.go @@ -235,8 +235,8 @@ func (t *Tracker) TrackTimeToFirstToken(dur time.Duration) error { return t.events.TrackMetric(timeToFirstToken, t.context, float64(dur.Milliseconds()), t.trackData) } -// TrackUsage tracks the token usage for a model evaluation. -func (t *Tracker) TrackUsage(usage TokenUsage) error { +// TrackTokens tracks the token usage for a model evaluation. +func (t *Tracker) TrackTokens(usage TokenUsage) error { if usage.Set() { t.tokens = ldcommon.Some(usage) } @@ -269,6 +269,13 @@ func (t *Tracker) TrackUsage(usage TokenUsage) error { return nil } +// TrackUsage tracks token usage. +// +// Deprecated: Use TrackTokens instead. +func (t *Tracker) TrackUsage(usage TokenUsage) error { + return t.TrackTokens(usage) +} + func measureDurationOfTask[T any, A any]( stopwatch Stopwatch, arg A, @@ -335,8 +342,8 @@ func (t *Tracker) TrackRequest(task func(c *Config) (ProviderResponse, error)) ( } if usage.Usage.Set() { - // TrackUsage logs errors. - _ = t.TrackUsage(usage.Usage) + // TrackTokens logs errors. + _ = t.TrackTokens(usage.Usage) } return usage, nil diff --git a/ldai/tracker_test.go b/ldai/tracker_test.go index 29bb71e4..836ff026 100644 --- a/ldai/tracker_test.go +++ b/ldai/tracker_test.go @@ -245,13 +245,13 @@ func TestTracker_TrackFeedback(t *testing.T) { assert.ElementsMatch(t, []trackEvent{expectedPositiveEvent, expectedNegativeEvent}, events.events) } -func TestTracker_TrackUsage(t *testing.T) { +func TestTracker_TrackTokens(t *testing.T) { t.Run("only one field set, only one event", func(t *testing.T) { events := newMockEvents() config := &Config{} tracker := newTracker("key", "variationKey", 8, events, config, ldcontext.New("key"), nil) - assert.NoError(t, tracker.TrackUsage(TokenUsage{ + assert.NoError(t, tracker.TrackTokens(TokenUsage{ Total: 42, })) @@ -270,7 +270,7 @@ func TestTracker_TrackUsage(t *testing.T) { config := &Config{} tracker := newTracker("key", "variationKey", 9, events, config, ldcontext.New("key"), nil) - assert.NoError(t, tracker.TrackUsage(TokenUsage{ + assert.NoError(t, tracker.TrackTokens(TokenUsage{ Total: 42, Input: 20, Output: 22, @@ -382,7 +382,7 @@ func TestTracker_GetSummary(t *testing.T) { Input: 40, Output: 60, } - _ = tracker.TrackUsage(usage) + _ = tracker.TrackTokens(usage) summary := tracker.GetSummary()