1package llm
2
3import "testing"
4
5func TestUsageTotalInputTokens(t *testing.T) {
6 tests := []struct {
7 name string
8 usage Usage
9 want uint64
10 }{
11 {
12 name: "all token types",
13 usage: Usage{
14 InputTokens: 100,
15 CacheCreationInputTokens: 50,
16 CacheReadInputTokens: 200,
17 OutputTokens: 30,
18 },
19 want: 350, // 100 + 50 + 200
20 },
21 {
22 name: "only input tokens",
23 usage: Usage{
24 InputTokens: 150,
25 OutputTokens: 50,
26 },
27 want: 150,
28 },
29 {
30 name: "heavy caching",
31 usage: Usage{
32 InputTokens: 10,
33 CacheCreationInputTokens: 0,
34 CacheReadInputTokens: 5000,
35 OutputTokens: 100,
36 },
37 want: 5010, // 10 + 0 + 5000
38 },
39 {
40 name: "zero",
41 usage: Usage{},
42 want: 0,
43 },
44 }
45
46 for _, tt := range tests {
47 t.Run(tt.name, func(t *testing.T) {
48 got := tt.usage.TotalInputTokens()
49 if got != tt.want {
50 t.Errorf("TotalInputTokens() = %d, want %d", got, tt.want)
51 }
52 })
53 }
54}
55
56func TestUsageContextWindowUsed(t *testing.T) {
57 tests := []struct {
58 name string
59 usage Usage
60 want uint64
61 }{
62 {
63 name: "all token types",
64 usage: Usage{
65 InputTokens: 100,
66 CacheCreationInputTokens: 50,
67 CacheReadInputTokens: 200,
68 OutputTokens: 30,
69 },
70 want: 380, // 100 + 50 + 200 + 30
71 },
72 {
73 name: "only input and output",
74 usage: Usage{
75 InputTokens: 150,
76 OutputTokens: 50,
77 },
78 want: 200,
79 },
80 {
81 name: "heavy caching with output",
82 usage: Usage{
83 InputTokens: 10,
84 CacheCreationInputTokens: 0,
85 CacheReadInputTokens: 5000,
86 OutputTokens: 100,
87 },
88 want: 5110, // 10 + 0 + 5000 + 100
89 },
90 }
91
92 for _, tt := range tests {
93 t.Run(tt.name, func(t *testing.T) {
94 got := tt.usage.ContextWindowUsed()
95 if got != tt.want {
96 t.Errorf("ContextWindowUsed() = %d, want %d", got, tt.want)
97 }
98 })
99 }
100}