127 lines
3.6 KiB
Go
127 lines
3.6 KiB
Go
package openai
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestEstimateInputTokens_SmallRequest(t *testing.T) {
|
|
// A small request with a short prompt should be far below 100k threshold.
|
|
estimate := estimateInputTokens("hi", nil, nil, nil)
|
|
if estimate >= 100000 {
|
|
t.Fatalf("small request estimate too large: %d", estimate)
|
|
}
|
|
if estimate <= 0 {
|
|
t.Fatalf("small request estimate should be positive: %d", estimate)
|
|
}
|
|
}
|
|
|
|
func TestEstimateInputTokens_LargePayloadExceedsThreshold(t *testing.T) {
|
|
// 150k characters should produce an estimate above 100k.
|
|
largePayload := make([]byte, 150000)
|
|
for i := range largePayload {
|
|
largePayload[i] = 'x'
|
|
}
|
|
prompt := string(largePayload)
|
|
estimate := estimateInputTokens(prompt, nil, nil, nil)
|
|
// estimate = runes/4 + runes/16, so 150k ASCII characters stay below 100k.
|
|
// We need roughly 320k+ runes to reach the threshold with this formula.
|
|
if estimate >= 100000 {
|
|
t.Fatalf("expected estimate below 100k for 150k chars: %d", estimate)
|
|
}
|
|
|
|
// 400k characters should produce above 100k.
|
|
largePayload2 := make([]byte, 400000)
|
|
for i := range largePayload2 {
|
|
largePayload2[i] = 'x'
|
|
}
|
|
prompt2 := string(largePayload2)
|
|
estimate2 := estimateInputTokens(prompt2, nil, nil, nil)
|
|
if estimate2 < 100000 {
|
|
t.Fatalf("expected estimate >= 100k for 400k chars: %d", estimate2)
|
|
}
|
|
}
|
|
|
|
func TestEstimateInputTokens_WithMetadata(t *testing.T) {
|
|
prompt := "test"
|
|
metadata := map[string]string{
|
|
"request_id": "req-001",
|
|
"workspace": "/home/user/project",
|
|
"custom_key": "custom_value_that_is_quite_long_for_testing",
|
|
}
|
|
estimate := estimateInputTokens(prompt, metadata, nil, nil)
|
|
if estimate <= 0 {
|
|
t.Fatalf("estimate with metadata should be positive: %d", estimate)
|
|
}
|
|
}
|
|
|
|
func TestEstimateInputTokens_WithTools(t *testing.T) {
|
|
prompt := "run the command"
|
|
tools := []any{
|
|
map[string]any{
|
|
"type": "function",
|
|
"function": map[string]any{
|
|
"name": "run_commands",
|
|
"parameters": map[string]any{
|
|
"type": "object",
|
|
"properties": map[string]any{
|
|
"commands": map[string]any{
|
|
"type": "array",
|
|
"items": map[string]any{"type": "string"},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
estimate := estimateInputTokens(prompt, nil, tools, nil)
|
|
if estimate <= 0 {
|
|
t.Fatalf("estimate with tools should be positive: %d", estimate)
|
|
}
|
|
}
|
|
|
|
func TestToolsSchemaToJSON_FloatLessThanOne(t *testing.T) {
|
|
got := toolsSchemaToJSON(map[string]any{
|
|
"type": "number",
|
|
"minimum": 0.5,
|
|
"multiple": 0.25,
|
|
})
|
|
if !strings.Contains(got, `"minimum":0.5`) {
|
|
t.Fatalf("expected minimum float to be encoded, got %q", got)
|
|
}
|
|
if !strings.Contains(got, `"multiple":0.25`) {
|
|
t.Fatalf("expected multiple float to be encoded, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestClassifyContext_Normal(t *testing.T) {
|
|
threshold := 100000
|
|
if classifyContext(1000, threshold) != "normal" {
|
|
t.Fatalf("expected normal for 1000 tokens")
|
|
}
|
|
if classifyContext(99999, threshold) != "normal" {
|
|
t.Fatalf("expected normal for 99999 tokens")
|
|
}
|
|
}
|
|
|
|
func TestClassifyContext_Long(t *testing.T) {
|
|
threshold := 100000
|
|
if classifyContext(100000, threshold) != "long" {
|
|
t.Fatalf("expected long for 100000 tokens")
|
|
}
|
|
if classifyContext(200000, threshold) != "long" {
|
|
t.Fatalf("expected long for 200000 tokens")
|
|
}
|
|
}
|
|
|
|
func TestClassifyContext_AtThreshold(t *testing.T) {
|
|
threshold := 100000
|
|
// At exactly the threshold, should be "long".
|
|
if classifyContext(threshold, threshold) != "long" {
|
|
t.Fatalf("expected long at threshold %d", threshold)
|
|
}
|
|
// Just below, should be "normal".
|
|
if classifyContext(threshold-1, threshold) != "normal" {
|
|
t.Fatalf("expected normal at threshold-1 %d", threshold-1)
|
|
}
|
|
}
|