loveuer
46 lines
1.3 KiB
Go
46 lines
1.3 KiB
Go
package llm
|
|
|
|
import "testing"
|
|
|
|
func TestEstimateTokensEmpty(t *testing.T) {
|
|
if got := EstimateTokens(nil, nil); got != 0 {
|
|
t.Fatalf("EstimateTokens(nil, nil) = %d, want 0", got)
|
|
}
|
|
}
|
|
|
|
func TestEstimateTokens(t *testing.T) {
|
|
shortMessages := []Message{{Role: RoleUser, Content: "hello"}}
|
|
short := EstimateTokens(shortMessages, nil)
|
|
if short <= 4 {
|
|
t.Fatalf("short estimate = %d, want message framing plus content", short)
|
|
}
|
|
if short > 50 {
|
|
t.Fatalf("short estimate = %d, looks too high for one short message", short)
|
|
}
|
|
|
|
longMessages := []Message{
|
|
{Role: RoleUser, Content: "hello"},
|
|
{Role: RoleAssistant, Content: "This is a longer response with enough content to exceed the short estimate."},
|
|
}
|
|
long := EstimateTokens(longMessages, nil)
|
|
if long <= short {
|
|
t.Fatalf("long estimate = %d, want > short estimate %d", long, short)
|
|
}
|
|
|
|
withTools := EstimateTokens(shortMessages, []Tool{{
|
|
Type: "function",
|
|
Function: ToolFunction{
|
|
Name: "file_read",
|
|
Description: "Read a file",
|
|
Parameters: []byte(`{"type":"object"}`),
|
|
},
|
|
}})
|
|
if withTools <= short {
|
|
t.Fatalf("withTools estimate = %d, want > short estimate %d", withTools, short)
|
|
}
|
|
|
|
if msgTokens := estimateMessageTokens(shortMessages[0]); msgTokens <= 4 {
|
|
t.Fatalf("estimateMessageTokens = %d, want message framing plus content", msgTokens)
|
|
}
|
|
}
|