forked from QuantumNous/new-api
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbilling_usage_test.go
More file actions
89 lines (74 loc) · 3.57 KB
/
Copy pathbilling_usage_test.go
File metadata and controls
89 lines (74 loc) · 3.57 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
package dto
import (
"testing"
kitutil "github.com/QuantumNous/new-api/relaykit/relayconvert/kitutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestNewGeminiChatBillingUsageRequiresTokenContent(t *testing.T) {
require.Nil(t, NewGeminiChatBillingUsage(nil))
require.Nil(t, NewGeminiChatBillingUsage(&GeminiUsageMetadata{}))
billingUsage := NewGeminiChatBillingUsage(&GeminiUsageMetadata{PromptTokenCount: 1})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.GeminiUsageMetadata)
assert.Equal(t, BillingUsageSourceGeminiChat, billingUsage.Source)
assert.Equal(t, BillingUsageSemanticGemini, billingUsage.Semantic)
assert.False(t, billingUsage.Estimated)
}
func TestNewClaudeMessagesBillingUsageRequiresTokenContent(t *testing.T) {
require.Nil(t, NewClaudeMessagesBillingUsage(nil))
require.Nil(t, NewClaudeMessagesBillingUsage(&ClaudeUsage{}))
require.Nil(t, NewClaudeMessagesBillingUsage(&ClaudeUsage{CacheCreation: &ClaudeCacheCreationUsage{}}))
billingUsage := NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 1})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.ClaudeUsage)
assert.Equal(t, BillingUsageSourceClaudeMessages, billingUsage.Source)
assert.Equal(t, BillingUsageSemanticAnthropic, billingUsage.Semantic)
cacheOnly := NewClaudeMessagesBillingUsage(&ClaudeUsage{
CacheCreation: &ClaudeCacheCreationUsage{Ephemeral5mInputTokens: 4},
})
require.NotNil(t, cacheOnly)
}
func TestNewOpenAIChatBillingUsageRequiresTokenContent(t *testing.T) {
require.Nil(t, NewOpenAIChatBillingUsage(nil))
require.Nil(t, NewOpenAIChatBillingUsage(&Usage{}))
billingUsage := NewOpenAIChatBillingUsage(&Usage{PromptTokens: 1})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.OpenAIUsage)
assert.Equal(t, BillingUsageSourceOAIChat, billingUsage.Source)
assert.Equal(t, BillingUsageSemanticOpenAI, billingUsage.Semantic)
assert.Equal(t, 1, billingUsage.OpenAIUsage.PromptTokens)
}
func TestNewEstimatedGeminiChatBillingUsage(t *testing.T) {
billingUsage := NewEstimatedGeminiChatBillingUsage(&Usage{
PromptTokens: 11,
CompletionTokens: 7,
})
require.NotNil(t, billingUsage)
require.NotNil(t, billingUsage.GeminiUsageMetadata)
assert.True(t, billingUsage.Estimated)
assert.Equal(t, 11, billingUsage.GeminiUsageMetadata.PromptTokenCount)
assert.Equal(t, 7, billingUsage.GeminiUsageMetadata.CandidatesTokenCount)
assert.Equal(t, 18, billingUsage.GeminiUsageMetadata.TotalTokenCount)
}
func TestBillingUsageJSONUsesProtocolNamedFields(t *testing.T) {
billingUsage := &BillingUsage{
OpenAIUsage: &Usage{PromptTokens: 1, BillingUsage: NewClaudeMessagesBillingUsage(&ClaudeUsage{InputTokens: 9})},
ClaudeUsage: &ClaudeUsage{InputTokens: 2, BillingUsage: NewOpenAIChatBillingUsage(&Usage{PromptTokens: 8})},
GeminiUsageMetadata: &GeminiUsageMetadata{PromptTokenCount: 3, BillingUsage: NewOpenAIChatBillingUsage(&Usage{PromptTokens: 7})},
}
data, err := kitutil.Marshal(billingUsage)
require.NoError(t, err)
assert.Contains(t, string(data), `"openai_usage"`)
assert.Contains(t, string(data), `"claude_usage"`)
assert.Contains(t, string(data), `"gemini_usage_metadata"`)
assert.NotContains(t, string(data), `"usage":`)
assert.NotContains(t, string(data), `"usage_metadata"`)
clone := CloneBillingUsage(billingUsage)
require.NotNil(t, clone.OpenAIUsage)
require.NotNil(t, clone.ClaudeUsage)
require.NotNil(t, clone.GeminiUsageMetadata)
assert.Nil(t, clone.OpenAIUsage.BillingUsage)
assert.Nil(t, clone.ClaudeUsage.BillingUsage)
assert.Nil(t, clone.GeminiUsageMetadata.BillingUsage)
}