You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

105 lines
3.2 KiB

  1. package service
  2. import (
  3. "fmt"
  4. "net/http/httptest"
  5. "testing"
  6. "github.com/QuantumNous/new-api/dto"
  7. "github.com/QuantumNous/new-api/types"
  8. "github.com/gin-gonic/gin"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) *gin.Context {
  12. rec := httptest.NewRecorder()
  13. ctx, _ := gin.CreateTestContext(rec)
  14. setChannelAffinityContext(ctx, channelAffinityMeta{
  15. CacheKey: fmt.Sprintf("test:%s:%s:%s", ruleName, usingGroup, keyFP),
  16. TTLSeconds: 600,
  17. RuleName: ruleName,
  18. UsingGroup: usingGroup,
  19. KeyFingerprint: keyFP,
  20. })
  21. return ctx
  22. }
  23. func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) {
  24. ruleName := "rule_" + t.Name()
  25. usingGroup := "default"
  26. keyFP := "fp_" + t.Name()
  27. ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
  28. usage := &dto.Usage{
  29. PromptTokens: 100,
  30. CompletionTokens: 40,
  31. TotalTokens: 140,
  32. PromptTokensDetails: dto.InputTokenDetails{
  33. CachedTokens: 30,
  34. },
  35. }
  36. ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, types.RelayFormatClaude)
  37. stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
  38. require.EqualValues(t, 1, stats.Total)
  39. require.EqualValues(t, 1, stats.Hit)
  40. require.EqualValues(t, 100, stats.PromptTokens)
  41. require.EqualValues(t, 40, stats.CompletionTokens)
  42. require.EqualValues(t, 140, stats.TotalTokens)
  43. require.EqualValues(t, 30, stats.CachedTokens)
  44. require.Equal(t, cacheTokenRateModeCachedOverPromptPlusCached, stats.CachedTokenRateMode)
  45. }
  46. func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) {
  47. ruleName := "rule_" + t.Name()
  48. usingGroup := "default"
  49. keyFP := "fp_" + t.Name()
  50. ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
  51. openAIUsage := &dto.Usage{
  52. PromptTokens: 100,
  53. PromptTokensDetails: dto.InputTokenDetails{
  54. CachedTokens: 10,
  55. },
  56. }
  57. claudeUsage := &dto.Usage{
  58. PromptTokens: 80,
  59. PromptTokensDetails: dto.InputTokenDetails{
  60. CachedTokens: 20,
  61. },
  62. }
  63. ObserveChannelAffinityUsageCacheByRelayFormat(ctx, openAIUsage, types.RelayFormatOpenAI)
  64. ObserveChannelAffinityUsageCacheByRelayFormat(ctx, claudeUsage, types.RelayFormatClaude)
  65. stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
  66. require.EqualValues(t, 2, stats.Total)
  67. require.EqualValues(t, 2, stats.Hit)
  68. require.EqualValues(t, 180, stats.PromptTokens)
  69. require.EqualValues(t, 30, stats.CachedTokens)
  70. require.Equal(t, cacheTokenRateModeMixed, stats.CachedTokenRateMode)
  71. }
  72. func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) {
  73. ruleName := "rule_" + t.Name()
  74. usingGroup := "default"
  75. keyFP := "fp_" + t.Name()
  76. ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
  77. usage := &dto.Usage{
  78. PromptTokens: 100,
  79. PromptTokensDetails: dto.InputTokenDetails{
  80. CachedTokens: 25,
  81. },
  82. }
  83. ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, types.RelayFormatGemini)
  84. stats := GetChannelAffinityUsageCacheStats(ruleName, usingGroup, keyFP)
  85. require.EqualValues(t, 1, stats.Total)
  86. require.EqualValues(t, 1, stats.Hit)
  87. require.EqualValues(t, 25, stats.CachedTokens)
  88. require.Equal(t, "", stats.CachedTokenRateMode)
  89. }