91 lines
2.4 KiB
Go
91 lines
2.4 KiB
Go
package ctxmgr
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestAssembleWindow_MaxTurnsBudget(t *testing.T) {
|
|
// 构造 20 条消息
|
|
var msgs []ChatMessageRecord
|
|
for i := 1; i <= 20; i++ {
|
|
role := "user"
|
|
if i%2 == 0 {
|
|
role = "assistant"
|
|
}
|
|
msgs = append(msgs, ChatMessageRecord{
|
|
ID: int64(i),
|
|
Role: role,
|
|
Content: fmt.Sprintf("这是测试消息第 %d 条内容", i),
|
|
})
|
|
}
|
|
|
|
cfg := WindowConfig{
|
|
MaxTurns: 8, // 最多 8 条
|
|
TokenBudget: 5000,
|
|
ReserveTail: 500,
|
|
}
|
|
|
|
window, early, truncated := AssembleWindow(msgs, cfg)
|
|
if !truncated {
|
|
t.Fatalf("expected truncated to be true")
|
|
}
|
|
if len(window) != 8 {
|
|
t.Errorf("expected 8 messages in window, got %d", len(window))
|
|
}
|
|
if len(early) != 12 {
|
|
t.Errorf("expected 12 messages in early, got %d", len(early))
|
|
}
|
|
// 检查窗口中第一条是否是第 13 条
|
|
if window[0].ID != 13 {
|
|
t.Errorf("expected window start with ID 13, got %d", window[0].ID)
|
|
}
|
|
// 检查窗口最后一条是否是第 20 条
|
|
if window[len(window)-1].ID != 20 {
|
|
t.Errorf("expected window end with ID 20, got %d", window[len(window)-1].ID)
|
|
}
|
|
}
|
|
|
|
func TestAssembleWindow_TokenBudgetLimit(t *testing.T) {
|
|
// 构造 5 条消息，每条内容长度 350 字符
|
|
var msgs []ChatMessageRecord
|
|
for i := 1; i <= 5; i++ {
|
|
msgs = append(msgs, ChatMessageRecord{
|
|
ID: int64(i),
|
|
Role: "user",
|
|
Content: strings.Repeat("长文本七字内容", 50), // 7 * 50 = 350 字符
|
|
})
|
|
}
|
|
|
|
cfg := WindowConfig{
|
|
MaxTurns: 10, // 条数未超限
|
|
TokenBudget: 1000, // 预算 1000
|
|
ReserveTail: 200, // 净可用 800，最多容纳 2 条 (300*2=600 <= 800)
|
|
}
|
|
|
|
window, early, truncated := AssembleWindow(msgs, cfg)
|
|
if !truncated {
|
|
t.Fatalf("expected truncated to be true")
|
|
}
|
|
if len(window) != 2 {
|
|
t.Errorf("expected 2 messages in window due to token budget, got %d", len(window))
|
|
}
|
|
if len(early) != 3 {
|
|
t.Errorf("expected 3 early messages, got %d", len(early))
|
|
}
|
|
}
|
|
|
|
func TestFormatNumberedDialogue(t *testing.T) {
|
|
msgs := []ChatMessageRecord{
|
|
{Role: "user", Content: "你好，请问你是谁？"},
|
|
{Role: "assistant", Content: "我是水木东方研发协同参谋。"},
|
|
}
|
|
|
|
formatted := FormatNumberedDialogue(msgs, 5)
|
|
expected := "[5] 用户: 你好，请问你是谁？\n[6] 助手: 我是水木东方研发协同参谋。\n"
|
|
if formatted != expected {
|
|
t.Errorf("expected formatted output:\n%s\ngot:\n%s", expected, formatted)
|
|
}
|
|
}
|