fix: follow live request results and prioritize usable groups

pkg/perf_metrics/live_status_test.go

@@ -2,35 +2,83 @@ package perfmetrics
 
 import (
 	"context"
-	relaycommon "github.com/QuantumNous/new-api/relay/common"
-	"github.com/QuantumNous/new-api/relaykit/types"
-	"github.com/stretchr/testify/assert"
 	"testing"
 	"time"
+
+	"github.com/QuantumNous/new-api/common"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/alicebob/miniredis/v2"
+	"github.com/go-redis/redis/v8"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
-func TestLiveUserSuccess(t *testing.T) {
-	t.Cleanup(func() {
-		for _, name := range []string{"user", "probe", "internal", "failed"} {
-			liveSuccess.Delete(name)
-		}
-	})
+func TestLiveUserResults(t *testing.T) {
+	oldRedis := common.RedisEnabled
+	common.RedisEnabled = false
+	t.Cleanup(func() { common.RedisEnabled = oldRedis })
+	canceled, cancel := context.WithCancel(context.Background())
+	cancel()
 	for _, tc := range []struct {
-		name  string
-		test  bool
-		token int
-		fail  bool
-		want  bool
+		name    string
+		test    bool
+		token   int
+		ctx     context.Context
+		err     *types.NewAPIError
+		want    bool
+		success bool
 	}{
-		{"user", false, 1, false, true}, {"probe", true, 1, false, false}, {"internal", false, 0, false, false}, {"failed", false, 1, true, false},
+		{"success", false, 1, context.Background(), nil, true, true},
+		{"upstream 503", false, 1, context.Background(), types.InitOpenAIError("server_error", 503), true, false},
+		{"probe", true, 1, context.Background(), types.InitOpenAIError("server_error", 503), false, false},
+		{"internal", false, 0, context.Background(), nil, false, false},
+		{"canceled", false, 1, canceled, types.InitOpenAIError("server_error", 503), false, false},
+		{"bad prompt", false, 1, context.Background(), types.InitOpenAIError("context_length_exceeded", 400), false, false},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			var apiErr *types.NewAPIError
-			if tc.fail {
-				apiErr = types.InitOpenAIError("server_error", 500)
+			t.Cleanup(func() { liveResults.Delete(tc.name) })
+			RecordRelayResult(tc.ctx, &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, tc.err)
+			result := LatestUserResults(context.Background(), []string{tc.name})[tc.name]
+			assert.Equal(t, tc.want, result.At > 0)
+			if tc.want {
+				assert.Equal(t, tc.success, result.Success)
 			}
-			RecordRelayResult(context.Background(), &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, apiErr)
-			assert.Equal(t, tc.want, LatestUserSuccess(context.Background(), []string{tc.name})[tc.name] > 0)
 		})
 	}
 }
+
+func TestUserResultsAreOrderedAndSurviveRestart(t *testing.T) {
+	server := miniredis.RunT(t)
+	client := redis.NewClient(&redis.Options{Addr: server.Addr()})
+	oldRedis, oldClient := common.RedisEnabled, common.RDB
+	common.RedisEnabled, common.RDB = true, client
+	t.Cleanup(func() {
+		common.RedisEnabled, common.RDB = oldRedis, oldClient
+		liveResults.Delete("group")
+		liveResults.Delete("other")
+		client.Close()
+	})
+	ctx := context.Background()
+	check := func(want UserRequestResult) {
+		assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
+		liveResults.Delete("group")
+		assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
+	}
+	recordUserResult("other", UserRequestResult{At: 2000, Success: true})
+	recordUserResult("group", UserRequestResult{At: 2000, Success: true})
+	check(UserRequestResult{At: 2000, Success: true})
+	recordUserResult("group", UserRequestResult{At: 2001, Success: false})
+	check(UserRequestResult{At: 2001, Success: false})
+	recordUserResult("group", UserRequestResult{At: 2000, Success: true})
+	check(UserRequestResult{At: 2001, Success: false})
+	recordUserResult("group", UserRequestResult{At: 2002, Success: true})
+	recordUserResult("group", UserRequestResult{At: 2001, Success: false})
+	check(UserRequestResult{At: 2002, Success: true})
+	assert.True(t, LatestUserResults(ctx, []string{"other"})["other"].Success)
+	ttl, err := client.TTL(ctx, "status:user-result:group").Result()
+	require.NoError(t, err)
+	assert.Equal(t, 15*time.Minute, ttl)
+	server.FastForward(15 * time.Minute)
+	assert.Zero(t, LatestUserResults(ctx, []string{"group"})["group"].At)
+}