134 lines
3.9 KiB
Go
134 lines
3.9 KiB
Go
package ice
|
|
|
|
import "testing"
|
|
|
|
func TestBudgetConfig_Calculate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
cfg BudgetConfig
|
|
promptTokens int
|
|
wantTotal int
|
|
wantConv int
|
|
wantMemory int
|
|
wantCode int
|
|
}{
|
|
{
|
|
name: "normal allocation",
|
|
cfg: BudgetConfig{
|
|
NumCtx: 8192,
|
|
SystemReserve: 1500,
|
|
RecentReserve: 2000,
|
|
ConversationPct: 0.40,
|
|
MemoryPct: 0.20,
|
|
CodePct: 0.40,
|
|
},
|
|
promptTokens: 500,
|
|
// available = int(8192*0.75) - 1500 - 2000 - 500 = 6144 - 4000 = 2144
|
|
wantTotal: 2144,
|
|
wantConv: 857, // int(2144 * 0.40) = 857
|
|
wantMemory: 428, // int(2144 * 0.20) = 428
|
|
wantCode: 857, // int(2144 * 0.40) = 857
|
|
},
|
|
{
|
|
name: "large prompt clamps to zero",
|
|
cfg: BudgetConfig{
|
|
NumCtx: 8192,
|
|
SystemReserve: 1500,
|
|
RecentReserve: 2000,
|
|
ConversationPct: 0.40,
|
|
MemoryPct: 0.20,
|
|
CodePct: 0.40,
|
|
},
|
|
promptTokens: 99999,
|
|
wantTotal: 0,
|
|
wantConv: 0,
|
|
wantMemory: 0,
|
|
wantCode: 0,
|
|
},
|
|
{
|
|
name: "exact boundary available is zero",
|
|
cfg: BudgetConfig{
|
|
NumCtx: 8192,
|
|
SystemReserve: 1500,
|
|
RecentReserve: 2000,
|
|
ConversationPct: 0.40,
|
|
MemoryPct: 0.20,
|
|
CodePct: 0.40,
|
|
},
|
|
// int(8192*0.75) - 1500 - 2000 = 2644
|
|
promptTokens: 2644,
|
|
wantTotal: 0,
|
|
wantConv: 0,
|
|
wantMemory: 0,
|
|
wantCode: 0,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
b := tt.cfg.Calculate(tt.promptTokens)
|
|
if b.Total != tt.wantTotal {
|
|
t.Errorf("Total = %d, want %d", b.Total, tt.wantTotal)
|
|
}
|
|
if b.Conversation != tt.wantConv {
|
|
t.Errorf("Conversation = %d, want %d", b.Conversation, tt.wantConv)
|
|
}
|
|
if b.Memory != tt.wantMemory {
|
|
t.Errorf("Memory = %d, want %d", b.Memory, tt.wantMemory)
|
|
}
|
|
if b.Code != tt.wantCode {
|
|
t.Errorf("Code = %d, want %d", b.Code, tt.wantCode)
|
|
}
|
|
if b.System != tt.cfg.SystemReserve {
|
|
t.Errorf("System = %d, want %d", b.System, tt.cfg.SystemReserve)
|
|
}
|
|
if b.Recent != tt.cfg.RecentReserve {
|
|
t.Errorf("Recent = %d, want %d", b.Recent, tt.cfg.RecentReserve)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestEstimateTokens(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
input string
|
|
want int
|
|
}{
|
|
{
|
|
name: "len/4 heuristic",
|
|
input: "hello world",
|
|
want: 2, // 11/4 = 2
|
|
},
|
|
{
|
|
name: "single char clamps to 1",
|
|
input: "a",
|
|
want: 1, // 1/4 = 0, clamp to 1
|
|
},
|
|
{
|
|
name: "empty string",
|
|
input: "",
|
|
want: 0,
|
|
},
|
|
{
|
|
name: "exactly 4 chars",
|
|
input: "abcd",
|
|
want: 1, // 4/4 = 1
|
|
},
|
|
{
|
|
name: "three chars clamps to 1",
|
|
input: "abc",
|
|
want: 1, // 3/4 = 0, clamp to 1
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got := estimateTokens(tt.input)
|
|
if got != tt.want {
|
|
t.Errorf("estimateTokens(%q) = %d, want %d", tt.input, got, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|