feat(agent): auto-compact long conversations in the runner
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.
This commit is contained in:
parent
4c5cad38f8
commit
7393812fba
2 changed files with 804 additions and 2 deletions
|
|
@ -5,12 +5,19 @@
|
||||||
// straight Q&A flow: no tools, no multi-iteration reasoning. Calling the
|
// straight Q&A flow: no tools, no multi-iteration reasoning. Calling the
|
||||||
// underlying provider keeps the prompt under our full control so we can
|
// underlying provider keeps the prompt under our full control so we can
|
||||||
// inject RAG context into the system message exactly where we want it.
|
// inject RAG context into the system message exactly where we want it.
|
||||||
|
//
|
||||||
|
// Auto-compaction (see CompactionConfig / Compact) folds the older portion
|
||||||
|
// of a long conversation into a single summary system message when the
|
||||||
|
// previous turn's input tokens approach the model's reported context
|
||||||
|
// window. It is opt-in via WithCompaction and short-circuits silently when
|
||||||
|
// the window is unknown or the history is too short to bother.
|
||||||
package agent
|
package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
"iter"
|
"iter"
|
||||||
|
"log/slog"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
||||||
|
|
@ -40,6 +47,9 @@ type Runner struct {
|
||||||
store *portfolio.Store
|
store *portfolio.Store
|
||||||
topK int
|
topK int
|
||||||
usage *Usage
|
usage *Usage
|
||||||
|
|
||||||
|
compaction CompactionConfig
|
||||||
|
lastCompact CompactionStats
|
||||||
}
|
}
|
||||||
|
|
||||||
type Usage struct {
|
type Usage struct {
|
||||||
|
|
@ -47,18 +57,64 @@ type Usage struct {
|
||||||
OutputTokens int
|
OutputTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// CompactionConfig mirrors the YAML `compaction:` block. Zero value means
|
||||||
|
// compaction is disabled.
|
||||||
|
type CompactionConfig struct {
|
||||||
|
Enabled bool
|
||||||
|
ThresholdRatio float64 // 0–1, default 0.75 when Enabled
|
||||||
|
KeepRecentTurns int // default 4 when Enabled
|
||||||
|
SummarySystemPrompt string // empty → built-in bilingual default
|
||||||
|
}
|
||||||
|
|
||||||
|
// CompactionStats reports the outcome of the most recent Compact call, so
|
||||||
|
// the HTTP handler can surface a "context compacted" event to the client.
|
||||||
|
type CompactionStats struct {
|
||||||
|
Happened bool // true when older turns were folded into a summary
|
||||||
|
OlderTurns int // user turns that got summarized away
|
||||||
|
KeptTurns int // user turns kept verbatim after the summary
|
||||||
|
SummaryTokens int // approx. token count of the summary text
|
||||||
|
WindowTokens int // model's reported context window at compaction time
|
||||||
|
UsedTokens int // input tokens reported by the previous turn
|
||||||
|
}
|
||||||
|
|
||||||
// New returns a Runner. The store may be nil (RAG disabled).
|
// New returns a Runner. The store may be nil (RAG disabled).
|
||||||
// systemPrompt is the full hand-written prompt from configs/...yaml.
|
// systemPrompt is the full hand-written prompt from configs/...yaml.
|
||||||
func New(client llm.LLMClient, p llmpersona.Persona, systemPrompt string, store *portfolio.Store, topK int) *Runner {
|
func New(client llm.LLMClient, p llmpersona.Persona, systemPrompt string, store *portfolio.Store, topK int) *Runner {
|
||||||
return &Runner{client: client, persona: p, systemPrompt: systemPrompt, store: store, topK: topK, usage: &Usage{}}
|
return &Runner{client: client, persona: p, systemPrompt: systemPrompt, store: store, topK: topK, usage: &Usage{}}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// WithCompaction enables auto-compaction with the given config. Returns
|
||||||
|
// the receiver for chaining. Pass a zero-value CompactionConfig to keep
|
||||||
|
// compaction disabled.
|
||||||
|
func (r *Runner) WithCompaction(cfg CompactionConfig) *Runner {
|
||||||
|
r.compaction = cfg
|
||||||
|
if cfg.Enabled {
|
||||||
|
if r.compaction.ThresholdRatio <= 0 || r.compaction.ThresholdRatio > 1 {
|
||||||
|
r.compaction.ThresholdRatio = 0.75
|
||||||
|
}
|
||||||
|
if r.compaction.KeepRecentTurns <= 0 {
|
||||||
|
r.compaction.KeepRecentTurns = 4
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return r
|
||||||
|
}
|
||||||
|
|
||||||
// LastUsage returns the token usage recorded on the most recent call.
|
// LastUsage returns the token usage recorded on the most recent call.
|
||||||
func (r *Runner) LastUsage() Usage { return *r.usage }
|
func (r *Runner) LastUsage() Usage { return *r.usage }
|
||||||
|
|
||||||
|
// LastCompaction returns the outcome of the most recent Compact call.
|
||||||
|
// Useful for the HTTP handler to emit a "compaction" SSE event after a
|
||||||
|
// stream finishes.
|
||||||
|
func (r *Runner) LastCompaction() CompactionStats { return r.lastCompact }
|
||||||
|
|
||||||
// BuildMessages prepares the system prompt and turns the chat history into
|
// BuildMessages prepares the system prompt and turns the chat history into
|
||||||
// the upstream message list. The system prompt includes RAG context for the
|
// the upstream message list. The system prompt includes RAG context for the
|
||||||
// user's last message (if RAG is enabled).
|
// user's last message (if RAG is enabled).
|
||||||
|
//
|
||||||
|
// Compaction is NOT applied here — callers invoke Compact explicitly before
|
||||||
|
// BuildMessages so the same Compact call doesn't run twice per request
|
||||||
|
// (Stream calls BuildMessages internally, and handlers sometimes call it
|
||||||
|
// first to extract the RAG context for the sources event).
|
||||||
func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Message, string, error) {
|
func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Message, string, error) {
|
||||||
ragContext := ""
|
ragContext := ""
|
||||||
if r.store != nil && len(history) > 0 {
|
if r.store != nil && len(history) > 0 {
|
||||||
|
|
@ -69,7 +125,7 @@ func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Messag
|
||||||
return nil, "", fmt.Errorf("rag search: %w", err)
|
return nil, "", fmt.Errorf("rag search: %w", err)
|
||||||
}
|
}
|
||||||
if len(hits) > 0 {
|
if len(hits) > 0 {
|
||||||
ragContext = formatHits(hits)
|
ragContext = r.limitRAGContext(formatHits(hits), history)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -77,12 +133,205 @@ func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Messag
|
||||||
// The system prompt comes from the YAML, not from the persona struct.
|
// The system prompt comes from the YAML, not from the persona struct.
|
||||||
// Keep the persona around only for the UI greeting.
|
// Keep the persona around only for the UI greeting.
|
||||||
system := botpersona.BuildSystemPrompt(r.systemPrompt, ragContext)
|
system := botpersona.BuildSystemPrompt(r.systemPrompt, ragContext)
|
||||||
|
history, err := r.fitHistory(system, history)
|
||||||
|
if err != nil {
|
||||||
|
return nil, "", err
|
||||||
|
}
|
||||||
msgs := make([]Message, 0, len(history)+1)
|
msgs := make([]Message, 0, len(history)+1)
|
||||||
msgs = append(msgs, Message{Role: RoleSystem, Content: system})
|
msgs = append(msgs, Message{Role: RoleSystem, Content: system})
|
||||||
msgs = append(msgs, history...)
|
msgs = append(msgs, history...)
|
||||||
return msgs, ragContext, nil
|
return msgs, ragContext, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const maxContextSafetyMargin = 256
|
||||||
|
|
||||||
|
func contextBudget(window int) int {
|
||||||
|
margin := int(float64(window) * 0.125)
|
||||||
|
if margin > maxContextSafetyMargin {
|
||||||
|
margin = maxContextSafetyMargin
|
||||||
|
}
|
||||||
|
if margin < 1 {
|
||||||
|
margin = 1
|
||||||
|
}
|
||||||
|
return window - margin
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Runner) limitRAGContext(ragContext string, history []Message) string {
|
||||||
|
if strings.TrimSpace(ragContext) == "" {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
budget := contextBudget(r.client.Capabilities().MaxContextWindow)
|
||||||
|
if budget <= 0 {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
|
||||||
|
historyTokens := totalPromptTokens(history)
|
||||||
|
fitted := ""
|
||||||
|
for _, block := range strings.Split(strings.TrimSpace(ragContext), "\n\n") {
|
||||||
|
block = strings.TrimSpace(block)
|
||||||
|
if block == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
candidate := block
|
||||||
|
if fitted != "" {
|
||||||
|
candidate = fitted + "\n\n" + block
|
||||||
|
}
|
||||||
|
if estimatePromptTokens(botpersona.BuildSystemPrompt(r.systemPrompt, candidate))+historyTokens > budget {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
fitted = candidate
|
||||||
|
}
|
||||||
|
return fitted
|
||||||
|
}
|
||||||
|
|
||||||
|
func (r *Runner) fitHistory(system string, history []Message) ([]Message, error) {
|
||||||
|
window := r.client.Capabilities().MaxContextWindow
|
||||||
|
if window <= 0 {
|
||||||
|
return history, nil
|
||||||
|
}
|
||||||
|
budget := contextBudget(window)
|
||||||
|
if budget <= 0 {
|
||||||
|
return nil, fmt.Errorf("context window is too small")
|
||||||
|
}
|
||||||
|
|
||||||
|
history = dropLeadingAssistantMessages(history)
|
||||||
|
for {
|
||||||
|
if estimatePromptTokens(system)+totalPromptTokens(history) <= budget {
|
||||||
|
return history, nil
|
||||||
|
}
|
||||||
|
turnStarts := []int{}
|
||||||
|
for i, message := range history {
|
||||||
|
if message.Role == RoleUser {
|
||||||
|
turnStarts = append(turnStarts, i)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(turnStarts) < 2 {
|
||||||
|
return nil, fmt.Errorf("prompt exceeds context window (%d tokens)", window)
|
||||||
|
}
|
||||||
|
first, next := turnStarts[0], turnStarts[1]
|
||||||
|
kept := make([]Message, 0, len(history)-(next-first))
|
||||||
|
kept = append(kept, history[:first]...)
|
||||||
|
kept = append(kept, history[next:]...)
|
||||||
|
history = kept
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func dropLeadingAssistantMessages(history []Message) []Message {
|
||||||
|
firstUser := -1
|
||||||
|
for i, message := range history {
|
||||||
|
if message.Role == RoleUser {
|
||||||
|
firstUser = i
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if firstUser <= 0 {
|
||||||
|
return history
|
||||||
|
}
|
||||||
|
out := make([]Message, 0, len(history))
|
||||||
|
for _, message := range history[:firstUser] {
|
||||||
|
if message.Role != RoleAssistant {
|
||||||
|
out = append(out, message)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return append(out, history[firstUser:]...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Compact reduces the older portion of history to a single system-role
|
||||||
|
// summary message when the previous turn's input tokens exceed the
|
||||||
|
// configured threshold. Returns the (possibly compacted) message slice —
|
||||||
|
// unchanged when compaction is disabled, history is too short, the model
|
||||||
|
// has no reported context window, or summarization fails.
|
||||||
|
//
|
||||||
|
// On summarization failure the runner falls back to dropping the oldest
|
||||||
|
// turns until the remaining slice fits a conservative budget, so a flaky
|
||||||
|
// summarize call never fails the user's request.
|
||||||
|
func (r *Runner) Compact(ctx context.Context, history []Message) ([]Message, error) {
|
||||||
|
r.lastCompact = CompactionStats{}
|
||||||
|
if !r.compaction.Enabled || len(history) == 0 {
|
||||||
|
return history, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
caps := r.client.Capabilities()
|
||||||
|
if caps.MaxContextWindow <= 0 {
|
||||||
|
return history, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
used := r.usage.InputTokens
|
||||||
|
estimated := estimatePromptTokens(r.systemPrompt) + totalPromptTokens(history)
|
||||||
|
if estimated > used {
|
||||||
|
used = estimated
|
||||||
|
}
|
||||||
|
window := caps.MaxContextWindow
|
||||||
|
threshold := int(float64(window) * r.compaction.ThresholdRatio)
|
||||||
|
if used < threshold {
|
||||||
|
return history, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
older, recent := splitByTurns(history, r.compaction.KeepRecentTurns)
|
||||||
|
if len(older) == 0 {
|
||||||
|
return history, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
summary, err := r.summarize(ctx, older)
|
||||||
|
if err != nil {
|
||||||
|
slog.Warn("compaction: summarize failed, falling back to truncation",
|
||||||
|
"err", err, "older_turns", userTurnCount(older))
|
||||||
|
// Fall back: drop oldest turns until the slice is small enough to
|
||||||
|
// fit in (threshold) tokens, using a 4-chars-per-token heuristic.
|
||||||
|
dropped, kept := truncateToBudget(history, threshold)
|
||||||
|
r.lastCompact = CompactionStats{
|
||||||
|
Happened: true,
|
||||||
|
OlderTurns: userTurnCount(dropped),
|
||||||
|
KeptTurns: userTurnCount(kept),
|
||||||
|
WindowTokens: window,
|
||||||
|
UsedTokens: used,
|
||||||
|
}
|
||||||
|
return kept, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
out := make([]Message, 0, 1+len(recent))
|
||||||
|
out = append(out, Message{
|
||||||
|
Role: RoleSystem,
|
||||||
|
Content: "Earlier conversation summary:\n" + summary,
|
||||||
|
})
|
||||||
|
out = append(out, recent...)
|
||||||
|
r.lastCompact = CompactionStats{
|
||||||
|
Happened: true,
|
||||||
|
OlderTurns: userTurnCount(older),
|
||||||
|
KeptTurns: userTurnCount(recent),
|
||||||
|
SummaryTokens: estimateTokens(summary),
|
||||||
|
WindowTokens: window,
|
||||||
|
UsedTokens: used,
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// summarize calls the LLM (non-streaming) to condense older turns. We use
|
||||||
|
// a small max_tokens cap so the compaction step itself stays cheap.
|
||||||
|
func (r *Runner) summarize(ctx context.Context, older []Message) (string, error) {
|
||||||
|
transcript := renderTranscript(older)
|
||||||
|
sys := r.compaction.SummarySystemPrompt
|
||||||
|
if sys == "" {
|
||||||
|
sys = defaultSummaryPrompt
|
||||||
|
}
|
||||||
|
maxTokens := 512
|
||||||
|
resp, err := r.client.Generate(ctx, llm.CompletionRequest{
|
||||||
|
Messages: []Message{
|
||||||
|
{Role: RoleSystem, Content: sys},
|
||||||
|
{Role: RoleUser, Content: transcript},
|
||||||
|
},
|
||||||
|
MaxTokens: &maxTokens,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
summary := strings.TrimSpace(resp.Content)
|
||||||
|
if summary == "" {
|
||||||
|
return "", fmt.Errorf("empty summary")
|
||||||
|
}
|
||||||
|
return summary, nil
|
||||||
|
}
|
||||||
|
|
||||||
// Stream runs the model and yields each streamed chunk. The caller
|
// Stream runs the model and yields each streamed chunk. The caller
|
||||||
// forwards chunk.Delta to the SSE stream; chunk.Usage on the last chunk
|
// forwards chunk.Delta to the SSE stream; chunk.Usage on the last chunk
|
||||||
// carries token counts.
|
// carries token counts.
|
||||||
|
|
@ -123,3 +372,131 @@ func formatHits(hits []portfolio.SearchResult) string {
|
||||||
}
|
}
|
||||||
return b.String()
|
return b.String()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// defaultSummaryPrompt is used when the YAML doesn't override it. The
|
||||||
|
// instruction explicitly avoids headers / meta-commentary so the summary
|
||||||
|
// drops in cleanly as a system message and the model treats it as facts
|
||||||
|
// about the prior conversation rather than instructions from the user.
|
||||||
|
const defaultSummaryPrompt = "Summarize the following conversation between a user and the assistant concisely but completely. " +
|
||||||
|
"Preserve decisions made, concrete facts (names, paths, IDs, preferences) and any pending task or context the assistant needs to keep helping. " +
|
||||||
|
"Respond only with the summary, in the same language as the conversation, with no headers or extra commentary."
|
||||||
|
|
||||||
|
// splitByTurns divides messages into an older portion to summarize and a
|
||||||
|
// recent tail to keep verbatim. A turn is a user message plus everything
|
||||||
|
// up to (but not including) the next user message — so a kept turn's tool
|
||||||
|
// calls, thoughts and response always stay together. Returns older == nil
|
||||||
|
// when there are not enough turns to bother compacting.
|
||||||
|
func splitByTurns(history []Message, keepRecentTurns int) (older, recent []Message) {
|
||||||
|
turnStarts := []int{}
|
||||||
|
for i, m := range history {
|
||||||
|
if m.Role == RoleUser {
|
||||||
|
turnStarts = append(turnStarts, i)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(turnStarts) <= keepRecentTurns {
|
||||||
|
return nil, history
|
||||||
|
}
|
||||||
|
cut := turnStarts[len(turnStarts)-keepRecentTurns]
|
||||||
|
return history[:cut], history[cut:]
|
||||||
|
}
|
||||||
|
|
||||||
|
// renderTranscript flattens chat messages into a plain User/Assistant
|
||||||
|
// transcript for summarization. Tool messages and empty assistant
|
||||||
|
// placeholders are skipped — neither contributes facts the model needs
|
||||||
|
// to remember.
|
||||||
|
func renderTranscript(messages []Message) string {
|
||||||
|
var b strings.Builder
|
||||||
|
for _, m := range messages {
|
||||||
|
switch m.Role {
|
||||||
|
case RoleUser:
|
||||||
|
b.WriteString("User: ")
|
||||||
|
b.WriteString(m.Content)
|
||||||
|
b.WriteString("\n")
|
||||||
|
case RoleAssistant:
|
||||||
|
if m.Content == "" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
b.WriteString("Assistant: ")
|
||||||
|
b.WriteString(m.Content)
|
||||||
|
b.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return b.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
// userTurnCount is the number of user-role messages in a slice — a
|
||||||
|
// convenient proxy for "turns" in stats reporting (each turn starts
|
||||||
|
// with a user message).
|
||||||
|
func userTurnCount(messages []Message) int {
|
||||||
|
n := 0
|
||||||
|
for _, m := range messages {
|
||||||
|
if m.Role == RoleUser {
|
||||||
|
n++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
|
||||||
|
// estimateTokens gives a rough token count for s. ~4 chars per token is
|
||||||
|
// the usual rule of thumb for English; close enough for Spanish / mixed
|
||||||
|
// text that this only drives a trigger threshold.
|
||||||
|
func estimateTokens(s string) int {
|
||||||
|
if s == "" {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return len(s)/4 + 1
|
||||||
|
}
|
||||||
|
|
||||||
|
func estimatePromptTokens(s string) int {
|
||||||
|
if s == "" {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return len(s)/3 + 1
|
||||||
|
}
|
||||||
|
|
||||||
|
func totalPromptTokens(messages []Message) int {
|
||||||
|
total := 0
|
||||||
|
for _, m := range messages {
|
||||||
|
total += estimatePromptTokens(m.Content)
|
||||||
|
}
|
||||||
|
return total
|
||||||
|
}
|
||||||
|
|
||||||
|
// truncateToBudget is the fallback when summarization fails. It drops
|
||||||
|
// the oldest user-turns one at a time until the remaining slice fits in
|
||||||
|
// `budget` tokens (conservative heuristic). At least the last user turn is
|
||||||
|
// always preserved, so we never return an empty slice.
|
||||||
|
//
|
||||||
|
// Implementation note: as `cut` increases, history[turnStarts[cut]:]
|
||||||
|
// shrinks, so `totalPromptTokens(...)` is monotonically non-increasing in
|
||||||
|
// `cut`. The loop is therefore guaranteed to terminate by either
|
||||||
|
// fitting the budget or hitting the "keep at least the last turn" floor.
|
||||||
|
func truncateToBudget(history []Message, budget int) (dropped, kept []Message) {
|
||||||
|
if len(history) == 0 {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
turnStarts := []int{}
|
||||||
|
for i, m := range history {
|
||||||
|
if m.Role == RoleUser {
|
||||||
|
turnStarts = append(turnStarts, i)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(turnStarts) == 0 {
|
||||||
|
return nil, history
|
||||||
|
}
|
||||||
|
// cut = number of leading turns to drop. 0 = keep everything; we
|
||||||
|
// grow it until the remaining slice fits in budget (or we hit the
|
||||||
|
// floor of one turn kept).
|
||||||
|
cut := 0
|
||||||
|
for cut < len(turnStarts)-1 {
|
||||||
|
start := turnStarts[cut+1]
|
||||||
|
if totalPromptTokens(history[start:]) <= budget {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
cut++
|
||||||
|
}
|
||||||
|
if cut == 0 {
|
||||||
|
return nil, history
|
||||||
|
}
|
||||||
|
return history[:turnStarts[cut]], history[turnStarts[cut]:]
|
||||||
|
}
|
||||||
|
|
|
||||||
425
internal/agent/runner_compaction_test.go
Normal file
425
internal/agent/runner_compaction_test.go
Normal file
|
|
@ -0,0 +1,425 @@
|
||||||
|
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))
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in a new issue