7393812fba
Runner.Compact folds the older portion of history into a single system-role summary when the previous turn's input tokens cross threshold_ratio × MaxContextWindow. When summarization fails, the runner falls back to truncateToBudget so a flaky summarize call never breaks the user's request. EstimatePromptTokens / totalPromptTokens give a conservative count (roughly 3 chars per token) used by both the compaction trigger and BuildMessages' new limitRAGContext / fitHistory helpers to cap the prompt inside the provider's reported window before the request goes out. Covers the first-turn case where no usage has been reported yet. Adds runner_compaction_test.go with table-driven coverage for the disabled, below-threshold, short-history, unknown-window, fallback and first-stream cases, plus a regression for the RAG-context trimmer.
425 lines
14 KiB
Go
425 lines
14 KiB
Go
package agent
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"iter"
|
||
"strings"
|
||
"testing"
|
||
|
||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||
llmpersona "github.com/VictorVargas/rony-llm-agent/pkg/persona"
|
||
)
|
||
|
||
// compactionStub is an llm.LLMClient that:
|
||
// - reports a configurable MaxContextWindow via Capabilities
|
||
// - records Generate calls (used for summarization)
|
||
// - returns a fixed Stream that also bumps the runner's r.usage via
|
||
// the usage chunk (so subsequent compactions see realistic usage)
|
||
type compactionStub struct {
|
||
window int
|
||
streamInput int // input tokens reported on the streamed usage chunk
|
||
generateCalls int
|
||
generateErr error
|
||
summary string
|
||
lastGenerate llm.CompletionRequest
|
||
}
|
||
|
||
func (s *compactionStub) Generate(_ context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
|
||
s.generateCalls++
|
||
s.lastGenerate = req
|
||
if s.generateErr != nil {
|
||
return llm.CompletionResponse{}, s.generateErr
|
||
}
|
||
return llm.CompletionResponse{
|
||
Content: s.summary,
|
||
StopReason: llm.StopReasonEndTurn,
|
||
}, nil
|
||
}
|
||
|
||
func (s *compactionStub) Stream(_ context.Context, _ llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
|
||
return func(yield func(llm.StreamChunk, error) bool) {
|
||
yield(llm.StreamChunk{Delta: "ok", FinishReason: "stop", Usage: llm.TokenUsage{InputTokens: s.streamInput, OutputTokens: 3}}, nil)
|
||
}
|
||
}
|
||
|
||
func (s *compactionStub) Name() string { return "compaction-stub" }
|
||
func (s *compactionStub) Capabilities() llm.ProviderCapabilities {
|
||
return llm.ProviderCapabilities{MaxContextWindow: s.window}
|
||
}
|
||
|
||
// pumpStream drives a single Stream call through the runner so the usage
|
||
// chunk updates r.usage — needed because the compaction trigger reads
|
||
// from r.usage, not from the request itself.
|
||
func pumpStream(t *testing.T, r *Runner, history []Message) {
|
||
t.Helper()
|
||
for chunk, err := range r.Stream(context.Background(), history) {
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
_ = chunk
|
||
}
|
||
}
|
||
|
||
// makeHistory builds a synthetic N-turn conversation (user + assistant
|
||
// pairs), each "long" enough to be obviously past any small threshold.
|
||
func makeHistory(turns int, contentLen int) []Message {
|
||
out := make([]Message, 0, turns*2)
|
||
payload := strings.Repeat("x", contentLen)
|
||
for i := 0; i < turns; i++ {
|
||
out = append(out,
|
||
Message{Role: RoleUser, Content: payload + " q"},
|
||
Message{Role: RoleAssistant, Content: payload + " a"},
|
||
)
|
||
}
|
||
return out
|
||
}
|
||
|
||
func personaMinimal() llmpersona.Persona {
|
||
return llmpersona.Persona{Name: "t", Tone: "concise", Language: "English"}
|
||
}
|
||
|
||
func TestCompactDisabledIsNoop(t *testing.T) {
|
||
cli := &compactionStub{window: 100, streamInput: 999, summary: "ignored"}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5) // no WithCompaction
|
||
pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
|
||
got, err := r.Compact(context.Background(), makeHistory(8, 200))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if cli.generateCalls != 0 {
|
||
t.Errorf("Generate called %d times, want 0 (compaction disabled)", cli.generateCalls)
|
||
}
|
||
if len(got) != 16 {
|
||
t.Errorf("len(history) = %d, want 16 (unchanged)", len(got))
|
||
}
|
||
stats := r.LastCompaction()
|
||
if stats.Happened {
|
||
t.Errorf("stats.Happened = true, want false")
|
||
}
|
||
}
|
||
|
||
func TestCompactBelowThresholdIsNoop(t *testing.T) {
|
||
// window = 1000, threshold = 0.5 → 500. Previous turn reported only
|
||
// 50 input tokens — well below the trigger.
|
||
cli := &compactionStub{window: 1000, streamInput: 50, summary: "ignored"}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5).
|
||
WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 2})
|
||
pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
|
||
got, err := r.Compact(context.Background(), makeHistory(8, 40))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if cli.generateCalls != 0 {
|
||
t.Errorf("Generate called %d times, want 0 (below threshold)", cli.generateCalls)
|
||
}
|
||
if len(got) != 16 {
|
||
t.Errorf("len(history) = %d, want 16 (unchanged)", len(got))
|
||
}
|
||
}
|
||
|
||
func TestCompactUsesEstimatedInputBeforeFirstStream(t *testing.T) {
|
||
cli := &compactionStub{window: 400, summary: "summary"}
|
||
r := New(cli, personaMinimal(), "system prompt", nil, 5).
|
||
WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 2})
|
||
|
||
got, err := r.Compact(context.Background(), makeHistory(8, 200))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if cli.generateCalls != 1 {
|
||
t.Fatalf("Generate called %d times, want 1", cli.generateCalls)
|
||
}
|
||
if len(got) != 5 {
|
||
t.Fatalf("len(compacted) = %d, want 5", len(got))
|
||
}
|
||
}
|
||
|
||
func TestCompactShortHistoryIsNoop(t *testing.T) {
|
||
// Above threshold but only 3 turns total (3 user, 3 assistant = 6
|
||
// messages). KeepRecentTurns = 4, so splitByTurns has nothing older
|
||
// to summarize.
|
||
cli := &compactionStub{window: 100, streamInput: 999, summary: "ignored"}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5).
|
||
WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 4})
|
||
pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
|
||
history := makeHistory(3, 200) // 6 messages, 3 user turns
|
||
got, err := r.Compact(context.Background(), history)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if cli.generateCalls != 0 {
|
||
t.Errorf("Generate called %d times, want 0 (not enough turns)", cli.generateCalls)
|
||
}
|
||
if len(got) != len(history) {
|
||
t.Errorf("len(history) = %d, want %d (unchanged)", len(got), len(history))
|
||
}
|
||
}
|
||
|
||
func TestCompactUnknownWindowIsNoop(t *testing.T) {
|
||
// Provider doesn't report a window (Capabilities returns 0).
|
||
cli := &compactionStub{window: 0, streamInput: 999, summary: "ignored"}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5).
|
||
WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 2})
|
||
pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
|
||
got, err := r.Compact(context.Background(), makeHistory(8, 200))
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if cli.generateCalls != 0 {
|
||
t.Errorf("Generate called %d times, want 0 (window unknown)", cli.generateCalls)
|
||
}
|
||
if len(got) != 16 {
|
||
t.Errorf("len(history) = %d, want 16 (unchanged)", len(got))
|
||
}
|
||
}
|
||
|
||
func TestCompactTriggersSummarize(t *testing.T) {
|
||
// 8 user turns × ~200 chars each = ~1600 chars ≈ 400 tokens. With
|
||
// window=400, threshold 0.5 = 200, and the stub reports 999 input
|
||
// tokens on the previous turn, compaction fires.
|
||
cli := &compactionStub{window: 400, streamInput: 999, summary: "user asked 8 things, here's the gist"}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5).
|
||
WithCompaction(CompactionConfig{
|
||
Enabled: true,
|
||
ThresholdRatio: 0.5,
|
||
KeepRecentTurns: 2,
|
||
})
|
||
pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
|
||
history := makeHistory(8, 200)
|
||
got, err := r.Compact(context.Background(), history)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if cli.generateCalls != 1 {
|
||
t.Fatalf("Generate called %d times, want 1 (compaction fired)", cli.generateCalls)
|
||
}
|
||
// The summary call's messages should be [system, user(transcript)].
|
||
if len(cli.lastGenerate.Messages) != 2 {
|
||
t.Fatalf("summarize call has %d messages, want 2", len(cli.lastGenerate.Messages))
|
||
}
|
||
if cli.lastGenerate.Messages[0].Role != RoleSystem {
|
||
t.Errorf("summarize call[0].Role = %q, want system", cli.lastGenerate.Messages[0].Role)
|
||
}
|
||
if cli.lastGenerate.Messages[1].Role != RoleUser {
|
||
t.Errorf("summarize call[1].Role = %q, want user", cli.lastGenerate.Messages[1].Role)
|
||
}
|
||
if !strings.Contains(cli.lastGenerate.Messages[1].Content, "User:") {
|
||
t.Errorf("summarize transcript missing 'User:' lines")
|
||
}
|
||
if got[0].Role != RoleSystem {
|
||
t.Errorf("compacted[0].Role = %q, want system (summary)", got[0].Role)
|
||
}
|
||
if !strings.Contains(got[0].Content, "user asked 8 things") {
|
||
t.Errorf("compacted[0] missing summary text: %q", got[0].Content)
|
||
}
|
||
// Recent tail: 2 turns = 4 messages, exactly the last 4 of the
|
||
// original history.
|
||
if len(got) != 1+4 {
|
||
t.Fatalf("len(compacted) = %d, want 5 (summary + 4)", len(got))
|
||
}
|
||
for i := 0; i < 4; i++ {
|
||
if got[1+i].Content != history[12+i].Content {
|
||
t.Errorf("kept[%d] content differs from history[%d]", i, 12+i)
|
||
}
|
||
}
|
||
stats := r.LastCompaction()
|
||
if !stats.Happened {
|
||
t.Fatal("stats.Happened = false, want true")
|
||
}
|
||
if stats.OlderTurns != 6 {
|
||
t.Errorf("OlderTurns = %d, want 6", stats.OlderTurns)
|
||
}
|
||
if stats.KeptTurns != 2 {
|
||
t.Errorf("KeptTurns = %d, want 2", stats.KeptTurns)
|
||
}
|
||
if stats.WindowTokens != 400 {
|
||
t.Errorf("WindowTokens = %d, want 400", stats.WindowTokens)
|
||
}
|
||
wantUsed := estimatePromptTokens("sys") + totalPromptTokens(history)
|
||
if stats.UsedTokens != wantUsed {
|
||
t.Errorf("UsedTokens = %d, want %d", stats.UsedTokens, wantUsed)
|
||
}
|
||
if stats.SummaryTokens == 0 {
|
||
t.Error("SummaryTokens = 0, want > 0")
|
||
}
|
||
}
|
||
|
||
func TestCompactSummarizeFailsFallsBackToTruncation(t *testing.T) {
|
||
cli := &compactionStub{
|
||
window: 400,
|
||
streamInput: 999,
|
||
generateErr: errors.New("summarize boom"),
|
||
}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5).
|
||
WithCompaction(CompactionConfig{
|
||
Enabled: true,
|
||
ThresholdRatio: 0.5,
|
||
KeepRecentTurns: 2,
|
||
})
|
||
pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
|
||
// 7 completed turns + 1 trailing user message (the request's new
|
||
// question). Real production requests always end on a user message.
|
||
history := append(makeHistory(7, 200), Message{Role: RoleUser, Content: "new question"})
|
||
got, err := r.Compact(context.Background(), history)
|
||
if err != nil {
|
||
t.Fatalf("Compact should swallow summarize failure, got err: %v", err)
|
||
}
|
||
if cli.generateCalls != 1 {
|
||
t.Errorf("Generate called %d times, want 1 (tried once before fallback)", cli.generateCalls)
|
||
}
|
||
// Fallback should have dropped the oldest turns. The trailing user
|
||
// message (the current question) must always be preserved.
|
||
last := got[len(got)-1]
|
||
if last.Role != RoleUser {
|
||
t.Errorf("last message role = %q, want user", last.Role)
|
||
}
|
||
if last.Content != "new question" {
|
||
t.Errorf("last message content = %q, want %q", last.Content, "new question")
|
||
}
|
||
stats := r.LastCompaction()
|
||
if !stats.Happened {
|
||
t.Error("stats.Happened = false, want true (fallback counts as compaction)")
|
||
}
|
||
if stats.OlderTurns == 0 {
|
||
t.Error("OlderTurns = 0, want > 0 (fallback dropped turns)")
|
||
}
|
||
}
|
||
|
||
func TestCompactEmptyHistoryIsNoop(t *testing.T) {
|
||
cli := &compactionStub{window: 100, streamInput: 999}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5).
|
||
WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.1, KeepRecentTurns: 2})
|
||
got, err := r.Compact(context.Background(), nil)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if got != nil {
|
||
t.Errorf("got = %v, want nil", got)
|
||
}
|
||
if cli.generateCalls != 0 {
|
||
t.Errorf("Generate called %d times, want 0", cli.generateCalls)
|
||
}
|
||
}
|
||
|
||
func TestFitHistoryDropsOldTurnsToFit(t *testing.T) {
|
||
cli := &compactionStub{window: 400}
|
||
r := New(cli, personaMinimal(), "sys", nil, 5)
|
||
history := makeHistory(4, 400)
|
||
|
||
got, err := r.fitHistory("sys", history)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if len(got) != 2 {
|
||
t.Fatalf("len(history) = %d, want 2", len(got))
|
||
}
|
||
if got[0].Role != RoleUser || got[1].Role != RoleAssistant {
|
||
t.Errorf("remaining turn roles = %q, %q", got[0].Role, got[1].Role)
|
||
}
|
||
}
|
||
|
||
func TestSplitByTurns(t *testing.T) {
|
||
history := []Message{
|
||
{Role: RoleUser, Content: "u1"},
|
||
{Role: RoleAssistant, Content: "a1"},
|
||
{Role: RoleUser, Content: "u2"},
|
||
{Role: RoleAssistant, Content: "a2"},
|
||
{Role: RoleUser, Content: "u3"},
|
||
{Role: RoleAssistant, Content: "a3"},
|
||
{Role: RoleUser, Content: "u4"},
|
||
{Role: RoleAssistant, Content: "a4"},
|
||
}
|
||
older, recent := splitByTurns(history, 2)
|
||
if len(older) != 4 {
|
||
t.Errorf("len(older) = %d, want 4", len(older))
|
||
}
|
||
if len(recent) != 4 {
|
||
t.Errorf("len(recent) = %d, want 4", len(recent))
|
||
}
|
||
if older[0].Content != "u1" || older[3].Content != "a2" {
|
||
t.Errorf("older = %v, want [u1,a1,u2,a2]", older)
|
||
}
|
||
if recent[0].Content != "u3" || recent[3].Content != "a4" {
|
||
t.Errorf("recent = %v, want [u3,a3,u4,a4]", recent)
|
||
}
|
||
}
|
||
|
||
func TestSplitByTurnsNothingToCompact(t *testing.T) {
|
||
history := []Message{
|
||
{Role: RoleUser, Content: "u1"},
|
||
{Role: RoleAssistant, Content: "a1"},
|
||
{Role: RoleUser, Content: "u2"},
|
||
}
|
||
older, recent := splitByTurns(history, 4)
|
||
if older != nil {
|
||
t.Errorf("older = %v, want nil", older)
|
||
}
|
||
if len(recent) != 3 {
|
||
t.Errorf("len(recent) = %d, want 3 (everything)", len(recent))
|
||
}
|
||
}
|
||
|
||
func TestRenderTranscriptSkipsToolAndEmpty(t *testing.T) {
|
||
msgs := []Message{
|
||
{Role: RoleUser, Content: "hi"},
|
||
{Role: RoleAssistant, Content: "hello!"},
|
||
{Role: RoleAssistant, Content: ""}, // streaming placeholder
|
||
{Role: RoleSystem, Content: "internal"},
|
||
}
|
||
got := renderTranscript(msgs)
|
||
want := "User: hi\nAssistant: hello!\n"
|
||
if got != want {
|
||
t.Errorf("renderTranscript =\n%q\nwant\n%q", got, want)
|
||
}
|
||
}
|
||
|
||
func TestEstimateTokens(t *testing.T) {
|
||
cases := []struct {
|
||
in string
|
||
want int
|
||
}{
|
||
{"", 0},
|
||
{"a", 1},
|
||
{"abcd", 2},
|
||
{strings.Repeat("x", 400), 101},
|
||
}
|
||
for _, c := range cases {
|
||
if got := estimateTokens(c.in); got != c.want {
|
||
t.Errorf("estimateTokens(%q) = %d, want %d", c.in, got, c.want)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestEstimatePromptTokensIsConservative(t *testing.T) {
|
||
text := strings.Repeat("x", 300)
|
||
if got := estimatePromptTokens(text); got <= estimateTokens(text) {
|
||
t.Errorf("estimatePromptTokens = %d, want greater than estimateTokens = %d", got, estimateTokens(text))
|
||
}
|
||
}
|
||
|
||
func TestTruncateToBudgetAlwaysKeepsLastUserTurn(t *testing.T) {
|
||
// 3 long completed turns + a short trailing user message (the new
|
||
// question). Budget so small only the last turn can fit.
|
||
long := strings.Repeat("x", 400)
|
||
history := []Message{
|
||
{Role: RoleUser, Content: long + " q1"},
|
||
{Role: RoleAssistant, Content: long + " a1"},
|
||
{Role: RoleUser, Content: long + " q2"},
|
||
{Role: RoleAssistant, Content: long + " a2"},
|
||
{Role: RoleUser, Content: long + " q3"},
|
||
{Role: RoleAssistant, Content: long + " a3"},
|
||
{Role: RoleUser, Content: "current question"},
|
||
}
|
||
dropped, kept := truncateToBudget(history, 10)
|
||
if len(kept) == 0 {
|
||
t.Fatal("kept is empty")
|
||
}
|
||
if kept[len(kept)-1].Content != "current question" {
|
||
t.Errorf("last kept message content = %q, want %q",
|
||
kept[len(kept)-1].Content, "current question")
|
||
}
|
||
if len(dropped)+len(kept) != len(history) {
|
||
t.Errorf("dropped+kept = %d, want %d", len(dropped)+len(kept), len(history))
|
||
}
|
||
}
|