feat: Add TrackError to mirror TrackSuccess (#225)

Additionally, emit new `$ld:ai:generation:(success|error)` events on success or failure.
launchdarkly · Dec 17, 2024 · ccd2c64 · ccd2c64
1 parent 345a3f0
commit ccd2c64
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 34 deletions.
diff --git a/ldai/tracker.go b/ldai/tracker.go
@@ -10,13 +10,15 @@ import (
 )
 
 const (
-	duration         = "$ld:ai:duration:total"
-	feedbackPositive = "$ld:ai:feedback:user:positive"
-	feedbackNegative = "$ld:ai:feedback:user:negative"
-	generation       = "$ld:ai:generation"
-	tokenTotal       = "$ld:ai:tokens:total"
-	tokenInput       = "$ld:ai:tokens:input"
-	tokenOutput      = "$ld:ai:tokens:output"
+	duration          = "$ld:ai:duration:total"
+	feedbackPositive  = "$ld:ai:feedback:user:positive"
+	feedbackNegative  = "$ld:ai:feedback:user:negative"
+	generation        = "$ld:ai:generation"
+	generationSuccess = "$ld:ai:generation:success"
+	generationError   = "$ld:ai:generation:error"
+	tokenTotal        = "$ld:ai:tokens:total"
+	tokenInput        = "$ld:ai:tokens:input"
+	tokenOutput       = "$ld:ai:tokens:output"
 )
 
 // TokenUsage represents the token usage returned by a model provider for a specific request.
@@ -167,7 +169,22 @@ func (t *Tracker) TrackFeedback(feedback Feedback) error {
 
 // TrackSuccess tracks a successful model evaluation.
 func (t *Tracker) TrackSuccess() error {
-	return t.events.TrackMetric(generation, t.context, 1, t.trackData)
+	err := t.events.TrackMetric(generation, t.context, 1, t.trackData)
+	if err := t.events.TrackMetric(generationSuccess, t.context, 1, t.trackData); err != nil {
+		return err
+	}
+
+	return err
+}
+
+// TrackError tracks an unsuccessful model evaluation.
+func (t *Tracker) TrackError() error {
+	err := t.events.TrackMetric(generation, t.context, 1, t.trackData)
+	if err := t.events.TrackMetric(generationError, t.context, 1, t.trackData); err != nil {
+		return err
+	}
+
+	return err
 }
 
 // TrackUsage tracks the token usage for a model evaluation.
@@ -222,6 +239,10 @@ func (t *Tracker) TrackRequest(task func(c *Config) (ProviderResponse, error)) (
 	usage, duration, err := measureDurationOfTask(t.stopwatch, t.config, task)
 
 	if err != nil {
+		if e := t.TrackError(); e != nil {
+			t.logWarning("error tracking error metric for request: %v", e)
+		}
+
 		t.logWarning("error executing request: %v", err)
 		return ProviderResponse{}, err
 	}

diff --git a/ldai/tracker_test.go b/ldai/tracker_test.go
@@ -57,14 +57,45 @@ func TestTracker_TrackSuccess(t *testing.T) {
 	tracker := newTracker("key", "variationKey", events, &Config{}, ldcontext.New("key"), nil)
 	assert.NoError(t, tracker.TrackSuccess())
 
-	expectedEvent := trackEvent{
-		name:        "$ld:ai:generation",
-		context:     ldcontext.New("key"),
-		metricValue: 1.0,
-		data:        makeTrackData("key", "variationKey"),
+	expectedEvents := []trackEvent{
+		{
+			name:        "$ld:ai:generation",
+			context:     ldcontext.New("key"),
+			metricValue: 1.0,
+			data:        makeTrackData("key", "variationKey"),
+		},
+		{
+			name:        "$ld:ai:generation:success",
+			context:     ldcontext.New("key"),
+			metricValue: 1.0,
+			data:        makeTrackData("key", "variationKey"),
+		},
 	}
 
-	assert.ElementsMatch(t, []trackEvent{expectedEvent}, events.events)
+	assert.ElementsMatch(t, expectedEvents, events.events)
+}
+
+func TestTracker_TrackError(t *testing.T) {
+	events := newMockEvents()
+	tracker := newTracker("key", "variationKey", events, &Config{}, ldcontext.New("key"), nil)
+	assert.NoError(t, tracker.TrackError())
+
+	expectedEvents := []trackEvent{
+		{
+			name:        "$ld:ai:generation",
+			context:     ldcontext.New("key"),
+			metricValue: 1.0,
+			data:        makeTrackData("key", "variationKey"),
+		},
+		{
+			name:        "$ld:ai:generation:error",
+			context:     ldcontext.New("key"),
+			metricValue: 1.0,
+			data:        makeTrackData("key", "variationKey"),
+		},
+	}
+
+	assert.ElementsMatch(t, expectedEvents, events.events)
 }
 
 func TestTracker_TrackRequest(t *testing.T) {
@@ -87,28 +118,33 @@ func TestTracker_TrackRequest(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, expectedResponse, r)
 
-	expectedSuccessEvent := trackEvent{
-		name:        "$ld:ai:generation",
-		context:     ldcontext.New("key"),
-		metricValue: 1,
-		data:        makeTrackData("key", "variationKey"),
-	}
-
-	expectedDurationEvent := trackEvent{
-		name:        "$ld:ai:duration:total",
-		context:     ldcontext.New("key"),
-		metricValue: 10.0,
-		data:        makeTrackData("key", "variationKey"),
-	}
-
-	expectedTokenUsageEvent := trackEvent{
-		name:        "$ld:ai:tokens:total",
-		context:     ldcontext.New("key"),
-		metricValue: 1,
-		data:        makeTrackData("key", "variationKey"),
+	expectedEvents := []trackEvent{
+		{
+			name:        "$ld:ai:generation",
+			context:     ldcontext.New("key"),
+			metricValue: 1,
+			data:        makeTrackData("key", "variationKey"),
+		},
+		{
+			name:        "$ld:ai:generation:success",
+			context:     ldcontext.New("key"),
+			metricValue: 1,
+			data:        makeTrackData("key", "variationKey"),
+		},
+		{
+			name:        "$ld:ai:duration:total",
+			context:     ldcontext.New("key"),
+			metricValue: 10.0,
+			data:        makeTrackData("key", "variationKey"),
+		},
+		{
+			name:        "$ld:ai:tokens:total",
+			context:     ldcontext.New("key"),
+			metricValue: 1,
+			data:        makeTrackData("key", "variationKey"),
+		},
 	}
 
-	expectedEvents := []trackEvent{expectedSuccessEvent, expectedDurationEvent, expectedTokenUsageEvent}
 	assert.ElementsMatch(t, expectedEvents, events.events)
 }