freedakgmail
131 lines
3.4 KiB
Go
131 lines
3.4 KiB
Go
package context
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/edgeai/gateway/internal/config"
|
|
"github.com/edgeai/gateway/pkg/api"
|
|
)
|
|
|
|
func TestTokenEstimator(t *testing.T) {
|
|
est := NewTokenEstimator()
|
|
|
|
// Empty string
|
|
if got := est.EstimateText(""); got != 0 {
|
|
t.Errorf("empty string: expected 0, got %d", got)
|
|
}
|
|
|
|
// English text
|
|
tokens := est.EstimateText("Hello world, this is a test.")
|
|
if tokens <= 0 {
|
|
t.Errorf("expected positive tokens for English, got %d", tokens)
|
|
}
|
|
|
|
// Chinese text (each char ~1 token)
|
|
cjkTokens := est.EstimateText("你好世界")
|
|
if cjkTokens != 4 {
|
|
t.Errorf("expected 4 tokens for 4 CJK chars, got %d", cjkTokens)
|
|
}
|
|
}
|
|
|
|
func TestEstimateKVCache(t *testing.T) {
|
|
// 1000 tokens, 32 layers, 4096 hidden dim, 2 bytes/element
|
|
result := EstimateKVCache(1000, 32, 4096, 2)
|
|
expected := int64(1000) * 32 * 2 * 4096 * 2
|
|
if result != expected {
|
|
t.Errorf("expected %d, got %d", expected, result)
|
|
}
|
|
}
|
|
|
|
func TestAssemblerNoTrim(t *testing.T) {
|
|
cfg := &config.ContextConfig{SafetyMarginRatio: 0.08}
|
|
a := NewAssembler(cfg)
|
|
|
|
history := []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: "Hello!"},
|
|
}
|
|
newMsgs := []api.Message{
|
|
{Role: "user", Content: "How are you?"},
|
|
}
|
|
|
|
result := a.Assemble(history, newMsgs, 1000, 100, "summary_and_recent")
|
|
if result.Trimmed {
|
|
t.Error("expected no trimming for small context")
|
|
}
|
|
if len(result.Messages) != 3 {
|
|
t.Errorf("expected 3 messages, got %d", len(result.Messages))
|
|
}
|
|
}
|
|
|
|
func TestAssemblerTrimRecentOnly(t *testing.T) {
|
|
cfg := &config.ContextConfig{SafetyMarginRatio: 0.08}
|
|
a := NewAssembler(cfg)
|
|
|
|
// Create many messages that exceed budget
|
|
msgs := make([]api.Message, 20)
|
|
for i := range msgs {
|
|
msgs[i] = api.Message{Role: "user", Content: "This is message number " + string(rune('A'+i))}
|
|
}
|
|
|
|
result := a.Assemble(msgs, []api.Message{}, 50, 10, "recent_only")
|
|
if !result.Trimmed {
|
|
t.Error("expected trimming for large context")
|
|
}
|
|
if len(result.Messages) >= 20 {
|
|
t.Error("expected fewer messages after trimming")
|
|
}
|
|
}
|
|
|
|
func TestAssemblerSummaryAndRecent(t *testing.T) {
|
|
cfg := &config.ContextConfig{SafetyMarginRatio: 0.08}
|
|
a := NewAssembler(cfg)
|
|
|
|
msgs := make([]api.Message, 0, 22)
|
|
msgs = append(msgs, api.Message{Role: "system", Content: "You are a helpful assistant."})
|
|
for i := 0; i < 20; i++ {
|
|
msgs = append(msgs, api.Message{Role: "user", Content: "Message " + string(rune('A'+i%26))})
|
|
msgs = append(msgs, api.Message{Role: "assistant", Content: "Response " + string(rune('A'+i%26))})
|
|
}
|
|
|
|
result := a.Assemble(msgs, []api.Message{}, 80, 20, "summary_and_recent")
|
|
if !result.Trimmed {
|
|
t.Error("expected trimming")
|
|
}
|
|
|
|
// System message should be preserved
|
|
hasSystem := false
|
|
hasSummary := false
|
|
for _, m := range result.Messages {
|
|
if m.Role == "system" {
|
|
if content, ok := m.Content.(string); ok {
|
|
if content == "You are a helpful assistant." {
|
|
hasSystem = true
|
|
}
|
|
if contains(content, "summarized") {
|
|
hasSummary = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if !hasSystem {
|
|
t.Error("system message should be preserved")
|
|
}
|
|
if !hasSummary {
|
|
t.Error("summary placeholder should be present when trimmed")
|
|
}
|
|
}
|
|
|
|
func contains(s, substr string) bool {
|
|
return len(s) >= len(substr) && (s == substr || (len(s) > len(substr) && (indexOf(s, substr) >= 0)))
|
|
}
|
|
|
|
func indexOf(s, substr string) int {
|
|
for i := 0; i <= len(s)-len(substr); i++ {
|
|
if s[i:i+len(substr)] == substr {
|
|
return i
|
|
}
|
|
}
|
|
return -1
|
|
}
|