120 lines
4.2 KiB
Go
120 lines
4.2 KiB
Go
package compaction
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/smallnest/pigo/internal/agentcore"
|
|
)
|
|
|
|
// bigUser builds a user message whose estimate is ~tokens tokens (4 chars each).
|
|
func bigUser(tokens int) agentcore.UserMessage {
|
|
return userMsg(strings.Repeat("x", tokens*4))
|
|
}
|
|
|
|
func assistantToolCall(id, name string) agentcore.AssistantMessage {
|
|
return agentcore.AssistantMessage{
|
|
RoleField: agentcore.RoleAssistant,
|
|
Content: agentcore.ContentList{agentcore.NewToolCallContent(id, name, nil)},
|
|
}
|
|
}
|
|
|
|
func toolResult(id string) agentcore.ToolResultMessage {
|
|
return agentcore.ToolResultMessage{
|
|
RoleField: agentcore.RoleToolResult,
|
|
ToolCallID: id,
|
|
Content: agentcore.ContentList{agentcore.NewTextContent("result")},
|
|
}
|
|
}
|
|
|
|
func TestFindCutPointNoValidCutPoint(t *testing.T) {
|
|
// Only toolResult messages -> no valid cut point.
|
|
msgs := []agentcore.Message{toolResult("a"), toolResult("b")}
|
|
got := FindCutPoint(msgs, 100)
|
|
if got.FirstKeptIndex != 0 || got.TurnStartIndex != -1 || got.IsSplitTurn {
|
|
t.Fatalf("no-cutpoint: got %+v, want {0 -1 false}", got)
|
|
}
|
|
}
|
|
|
|
func TestFindCutPointOnTurnBoundary(t *testing.T) {
|
|
// Three complete user/assistant turns, each user ~100 tokens.
|
|
// keepRecentTokens small so the cut lands on a clean user boundary.
|
|
msgs := []agentcore.Message{
|
|
bigUser(100), // 0
|
|
assistantMsg("a1", nil, ""), // 1
|
|
bigUser(100), // 2
|
|
assistantMsg("a2", nil, ""), // 3
|
|
bigUser(100), // 4
|
|
assistantMsg("a3", nil, ""), // 5
|
|
}
|
|
// keepRecentTokens=50: walking back, msg[5] tiny, msg[4]=100 >= 50 at i=4.
|
|
// nearest cut point >= 4 is index 4 (a user message) -> clean boundary.
|
|
got := FindCutPoint(msgs, 50)
|
|
if got.FirstKeptIndex != 4 {
|
|
t.Fatalf("FirstKeptIndex: got %d, want 4", got.FirstKeptIndex)
|
|
}
|
|
if got.IsSplitTurn {
|
|
t.Fatalf("IsSplitTurn: got true, want false (landed on user boundary)")
|
|
}
|
|
if got.TurnStartIndex != -1 {
|
|
t.Fatalf("TurnStartIndex: got %d, want -1", got.TurnStartIndex)
|
|
}
|
|
}
|
|
|
|
func TestFindCutPointSplitTurn(t *testing.T) {
|
|
// A turn starting with a user message, a big assistant reply, a toolCall and
|
|
// its toolResult. If the budget forces the cut onto the assistant message
|
|
// mid-turn, it's a split turn.
|
|
msgs := []agentcore.Message{
|
|
userMsg("older turn"), // 0
|
|
assistantMsg("older reply", nil, ""), // 1
|
|
userMsg("current turn start"), // 2 user
|
|
assistantMsg(strings.Repeat("y", 400), nil, ""), // 3 assistant ~100 tokens
|
|
assistantToolCall("t1", "read"), // 4 assistant (valid cut)
|
|
toolResult("t1"), // 5 toolResult (not cuttable)
|
|
}
|
|
// keepRecentTokens=50: from end, msg[5]=~2, msg[4]~1, msg[3]=100 >= 50 at i=3.
|
|
// nearest cut point >= 3 is index 3 (assistant) -> split turn, turn start=2.
|
|
got := FindCutPoint(msgs, 50)
|
|
if got.FirstKeptIndex != 3 {
|
|
t.Fatalf("FirstKeptIndex: got %d, want 3", got.FirstKeptIndex)
|
|
}
|
|
if !got.IsSplitTurn {
|
|
t.Fatalf("IsSplitTurn: got false, want true (cut on assistant mid-turn)")
|
|
}
|
|
if got.TurnStartIndex != 2 {
|
|
t.Fatalf("TurnStartIndex: got %d, want 2", got.TurnStartIndex)
|
|
}
|
|
}
|
|
|
|
func TestFindCutPointNeverCutsOnToolResult(t *testing.T) {
|
|
// Ensure a toolResult is never chosen even when it's the message where the
|
|
// budget is reached.
|
|
msgs := []agentcore.Message{
|
|
userMsg("u0"), // 0
|
|
assistantToolCall("t1", "grep"), // 1
|
|
toolResult("t1"), // 2 (budget could land here)
|
|
assistantMsg("done", nil, ""), // 3
|
|
}
|
|
got := FindCutPoint(msgs, 1) // tiny budget: reached at msg[3]
|
|
// cut must be a valid point (>=3 is index 3 assistant); never index 2.
|
|
if got.FirstKeptIndex == 2 {
|
|
t.Fatalf("cut landed on toolResult (index 2), which is illegal")
|
|
}
|
|
if msgs[got.FirstKeptIndex].Role() == agentcore.RoleToolResult {
|
|
t.Fatalf("FirstKeptIndex points at a toolResult: %d", got.FirstKeptIndex)
|
|
}
|
|
}
|
|
|
|
func TestFindCutPointBudgetNeverReachedKeepsEarliest(t *testing.T) {
|
|
msgs := []agentcore.Message{
|
|
userMsg("u0"), // 0
|
|
assistantMsg("a0", nil, ""), // 1
|
|
}
|
|
// Huge budget -> never reached -> keep from earliest valid cut point (0).
|
|
got := FindCutPoint(msgs, 1_000_000)
|
|
if got.FirstKeptIndex != 0 {
|
|
t.Fatalf("FirstKeptIndex: got %d, want 0 (earliest cut point)", got.FirstKeptIndex)
|
|
}
|
|
}
|