mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-03-30 02:05:21 +00:00
a01a77fc6f
* fix: claude affinity cache counter * fix: claude affinity cache counter * fix: stabilize cache usage stats format and simplify modal rendering
106 lines
3.4 KiB
Go
106 lines
3.4 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/types"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) *gin.Context {
|
|
rec := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(rec)
|
|
setChannelAffinityContext(ctx, channelAffinityMeta{
|
|
CacheKey: fmt.Sprintf("test:%s:%s:%s", ruleName, usingGroup, keyFP),
|
|
TTLSeconds: 600,
|
|
RuleName: ruleName,
|
|
UsingGroup: usingGroup,
|
|
KeyFingerprint: keyFP,
|
|
})
|
|
return ctx
|
|
}
|
|
|
|
func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) {
|
|
ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
|
|
usingGroup := "default"
|
|
keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
|
|
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
CompletionTokens: 40,
|
|
TotalTokens: 140,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 30,
|
|
},
|
|
}
|
|
|
|
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, types.RelayFormatClaude)
|
|
stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
|
|
|
|
require.EqualValues(t, 1, stats.Total)
|
|
require.EqualValues(t, 1, stats.Hit)
|
|
require.EqualValues(t, 100, stats.PromptTokens)
|
|
require.EqualValues(t, 40, stats.CompletionTokens)
|
|
require.EqualValues(t, 140, stats.TotalTokens)
|
|
require.EqualValues(t, 30, stats.CachedTokens)
|
|
require.Equal(t, cacheTokenRateModeCachedOverPromptPlusCached, stats.CachedTokenRateMode)
|
|
}
|
|
|
|
func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) {
|
|
ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
|
|
usingGroup := "default"
|
|
keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
|
|
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
|
|
|
|
openAIUsage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 10,
|
|
},
|
|
}
|
|
claudeUsage := &dto.Usage{
|
|
PromptTokens: 80,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 20,
|
|
},
|
|
}
|
|
|
|
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, openAIUsage, types.RelayFormatOpenAI)
|
|
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, claudeUsage, types.RelayFormatClaude)
|
|
stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
|
|
|
|
require.EqualValues(t, 2, stats.Total)
|
|
require.EqualValues(t, 2, stats.Hit)
|
|
require.EqualValues(t, 180, stats.PromptTokens)
|
|
require.EqualValues(t, 30, stats.CachedTokens)
|
|
require.Equal(t, cacheTokenRateModeMixed, stats.CachedTokenRateMode)
|
|
}
|
|
|
|
func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) {
|
|
ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
|
|
usingGroup := "default"
|
|
keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
|
|
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
|
|
|
|
usage := &dto.Usage{
|
|
PromptTokens: 100,
|
|
PromptTokensDetails: dto.InputTokenDetails{
|
|
CachedTokens: 25,
|
|
},
|
|
}
|
|
|
|
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, types.RelayFormatGemini)
|
|
stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
|
|
|
|
require.EqualValues(t, 1, stats.Total)
|
|
require.EqualValues(t, 1, stats.Hit)
|
|
require.EqualValues(t, 25, stats.CachedTokens)
|
|
require.Equal(t, "", stats.CachedTokenRateMode)
|
|
}
|