agent: add bounded context builder to prevent prompt explosion
Add agent/context.go with a ContextBuilder that enforces a character budget (soft + hard) over the assembled history, evicting oldest non-system messages first and replacing large tool outputs with truncated summaries. Update agent/session.go to use ContextBuilder in History(), so the loop always sees a bounded context window regardless of how many steps have run. The budget defaults are: - SoftLimit: 24000 chars (~6k tokens at 4 chars/token) - HardLimit: 96000 chars (~24k tokens) - MaxToolOutputChars: 2000 - TailMessages: 6 (always keep last N messages verbatim) All limits are configurable via ContextConfig.
This commit is contained in:
Levi Neely
parent
71ff058b1a
commit
833f9e0a00
|
|
@ -0,0 +1,296 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"ollie/backend"
|
||||
)
|
||||
|
||||
// ContextConfig controls the bounded context window behaviour.
|
||||
// All sizes are in characters (a rough proxy for tokens; ~4 chars per token).
|
||||
type ContextConfig struct {
|
||||
// SoftLimit: if assembled history exceeds this, begin evicting old messages.
|
||||
// Defaults to 24000 (~6k tokens).
|
||||
SoftLimit int
|
||||
|
||||
// HardLimit: absolute ceiling; messages are truncated to fit.
|
||||
// Defaults to 96000 (~24k tokens).
|
||||
HardLimit int
|
||||
|
||||
// MaxToolOutputChars: tool result messages longer than this are truncated
|
||||
// before being added to history. Defaults to 2000.
|
||||
MaxToolOutputChars int
|
||||
|
||||
// TailMessages: always preserve the most recent N messages verbatim,
|
||||
// even under eviction pressure. Defaults to 6.
|
||||
TailMessages int
|
||||
}
|
||||
|
||||
func defaultContextConfig() ContextConfig {
|
||||
return ContextConfig{
|
||||
SoftLimit: 24_000,
|
||||
HardLimit: 96_000,
|
||||
MaxToolOutputChars: 2_000,
|
||||
TailMessages: 6,
|
||||
}
|
||||
}
|
||||
|
||||
func (c *ContextConfig) setDefaults() {
|
||||
if c.SoftLimit <= 0 {
|
||||
c.SoftLimit = 24_000
|
||||
}
|
||||
if c.HardLimit <= 0 {
|
||||
c.HardLimit = 96_000
|
||||
}
|
||||
if c.MaxToolOutputChars <= 0 {
|
||||
c.MaxToolOutputChars = 2_000
|
||||
}
|
||||
if c.TailMessages <= 0 {
|
||||
c.TailMessages = 6
|
||||
}
|
||||
}
|
||||
|
||||
// ContextBuilder manages a rolling bounded history window.
|
||||
type ContextBuilder struct {
|
||||
cfg ContextConfig
|
||||
messages []backend.Message // full unbounded log
|
||||
}
|
||||
|
||||
// NewContextBuilder creates a ContextBuilder with the given config.
|
||||
// Pass a zero-value ContextConfig to use all defaults.
|
||||
func NewContextBuilder(cfg ContextConfig) *ContextBuilder {
|
||||
cfg.setDefaults()
|
||||
return &ContextBuilder{cfg: cfg}
|
||||
}
|
||||
|
||||
// Append adds a message to the full history.
|
||||
// Tool result messages are truncated to MaxToolOutputChars before storage.
|
||||
func (cb *ContextBuilder) Append(m backend.Message) {
|
||||
if m.Role == "tool" && len(m.Content) > cb.cfg.MaxToolOutputChars {
|
||||
m.Content = m.Content[:cb.cfg.MaxToolOutputChars] +
|
||||
fmt.Sprintf("\n... [truncated %d chars]", len(m.Content)-cb.cfg.MaxToolOutputChars)
|
||||
}
|
||||
cb.messages = append(cb.messages, m)
|
||||
}
|
||||
|
||||
// Messages returns the full stored message log (unbounded).
|
||||
func (cb *ContextBuilder) Messages() []backend.Message {
|
||||
return cb.messages
|
||||
}
|
||||
|
||||
// BoundedHistory returns a context-window-safe slice of messages.
|
||||
//
|
||||
// Strategy:
|
||||
// 1. System messages are always included at the front.
|
||||
// 2. The most recent TailMessages non-system messages are always kept.
|
||||
// 3. Older messages are included newest-first until SoftLimit is reached.
|
||||
// 4. If total still exceeds HardLimit, oldest non-system/non-tail messages
|
||||
// are dropped entirely.
|
||||
func (cb *ContextBuilder) BoundedHistory() []backend.Message {
|
||||
var system []backend.Message
|
||||
var rest []backend.Message
|
||||
|
||||
for _, m := range cb.messages {
|
||||
if m.Role == "system" {
|
||||
system = append(system, m)
|
||||
} else {
|
||||
rest = append(rest, m)
|
||||
}
|
||||
}
|
||||
|
||||
if len(rest) == 0 {
|
||||
return system
|
||||
}
|
||||
|
||||
// Always keep the tail.
|
||||
tailStart := len(rest) - cb.cfg.TailMessages
|
||||
if tailStart < 0 {
|
||||
tailStart = 0
|
||||
}
|
||||
tail := rest[tailStart:]
|
||||
older := rest[:tailStart]
|
||||
|
||||
// Budget remaining after system + tail.
|
||||
used := msgSliceChars(system) + msgSliceChars(tail)
|
||||
budget := cb.cfg.SoftLimit - used
|
||||
|
||||
// Greedily include older messages newest-first until budget exhausted.
|
||||
var included []backend.Message
|
||||
for i := len(older) - 1; i >= 0; i-- {
|
||||
size := msgChars(older[i])
|
||||
if budget-size < 0 {
|
||||
break
|
||||
}
|
||||
budget -= size
|
||||
included = append([]backend.Message{older[i]}, included...)
|
||||
}
|
||||
|
||||
result := make([]backend.Message, 0, len(system)+len(included)+len(tail))
|
||||
result = append(result, system...)
|
||||
result = append(result, included...)
|
||||
result = append(result, tail...)
|
||||
|
||||
// Hard-limit safety: truncate from the front (after system) if still over.
|
||||
for totalChars(result) > cb.cfg.HardLimit && len(result) > len(system)+1 {
|
||||
// Drop the oldest non-system message.
|
||||
result = append(result[:len(system)], result[len(system)+1:]...)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// Len returns the number of stored messages.
|
||||
func (cb *ContextBuilder) Len() int { return len(cb.messages) }
|
||||
|
||||
// ApproxTokens returns a rough token estimate for the bounded history
|
||||
// using the 4-chars-per-token heuristic.
|
||||
func (cb *ContextBuilder) ApproxTokens() int {
|
||||
return totalChars(cb.BoundedHistory()) / 4
|
||||
}
|
||||
|
||||
// --- helpers ----------------------------------------------------------------
|
||||
|
||||
func msgChars(m backend.Message) int {
|
||||
n := len(m.Content)
|
||||
for _, tc := range m.ToolCalls {
|
||||
n += len(tc.Name) + len(tc.Arguments)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func msgSliceChars(msgs []backend.Message) int {
|
||||
total := 0
|
||||
for _, m := range msgs {
|
||||
total += msgChars(m)
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
func totalChars(msgs []backend.Message) int {
|
||||
return msgSliceChars(msgs)
|
||||
}
|
||||
|
||||
// contextSummaryLine produces a one-line summary injected when old messages
|
||||
// are evicted, so the model knows compaction occurred.
|
||||
func contextSummaryLine(evicted int) backend.Message {
|
||||
return backend.Message{
|
||||
Role: "user",
|
||||
Content: fmt.Sprintf("[context compacted: %d earlier messages omitted to stay within token budget]", evicted),
|
||||
}
|
||||
}
|
||||
|
||||
// BoundedHistoryWithNotice is like BoundedHistory but injects a compaction
|
||||
// notice message when older messages were dropped, so the model is aware.
|
||||
func (cb *ContextBuilder) BoundedHistoryWithNotice() []backend.Message {
|
||||
var system []backend.Message
|
||||
var rest []backend.Message
|
||||
|
||||
for _, m := range cb.messages {
|
||||
if m.Role == "system" {
|
||||
system = append(system, m)
|
||||
} else {
|
||||
rest = append(rest, m)
|
||||
}
|
||||
}
|
||||
|
||||
if len(rest) == 0 {
|
||||
return system
|
||||
}
|
||||
|
||||
tailStart := len(rest) - cb.cfg.TailMessages
|
||||
if tailStart < 0 {
|
||||
tailStart = 0
|
||||
}
|
||||
tail := rest[tailStart:]
|
||||
older := rest[:tailStart]
|
||||
|
||||
used := msgSliceChars(system) + msgSliceChars(tail)
|
||||
budget := cb.cfg.SoftLimit - used
|
||||
|
||||
var included []backend.Message
|
||||
for i := len(older) - 1; i >= 0; i-- {
|
||||
size := msgChars(older[i])
|
||||
if budget-size < 0 {
|
||||
break
|
||||
}
|
||||
budget -= size
|
||||
included = append([]backend.Message{older[i]}, included...)
|
||||
}
|
||||
|
||||
evicted := len(older) - len(included)
|
||||
|
||||
result := make([]backend.Message, 0, len(system)+len(included)+len(tail)+1)
|
||||
result = append(result, system...)
|
||||
if evicted > 0 {
|
||||
result = append(result, contextSummaryLine(evicted))
|
||||
}
|
||||
result = append(result, included...)
|
||||
result = append(result, tail...)
|
||||
|
||||
for totalChars(result) > cb.cfg.HardLimit && len(result) > len(system)+1 {
|
||||
result = append(result[:len(system)], result[len(system)+1:]...)
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// ContextStats describes the current state of the context window.
|
||||
type ContextStats struct {
|
||||
StoredMessages int
|
||||
BoundedMessages int
|
||||
ApproxTokens int
|
||||
Evicted int
|
||||
SoftLimit int
|
||||
HardLimit int
|
||||
}
|
||||
|
||||
func (cb *ContextBuilder) Stats() ContextStats {
|
||||
bounded := cb.BoundedHistory()
|
||||
var system []backend.Message
|
||||
for _, m := range cb.messages {
|
||||
if m.Role == "system" {
|
||||
system = append(system, m)
|
||||
}
|
||||
}
|
||||
evicted := len(cb.messages) - len(bounded)
|
||||
if evicted < 0 {
|
||||
evicted = 0
|
||||
}
|
||||
return ContextStats{
|
||||
StoredMessages: len(cb.messages),
|
||||
BoundedMessages: len(bounded),
|
||||
ApproxTokens: totalChars(bounded) / 4,
|
||||
Evicted: evicted,
|
||||
SoftLimit: cb.cfg.SoftLimit,
|
||||
HardLimit: cb.cfg.HardLimit,
|
||||
}
|
||||
}
|
||||
|
||||
// ContextStatsString returns a one-line human-readable summary.
|
||||
func (cb *ContextBuilder) ContextStatsString() string {
|
||||
s := cb.Stats()
|
||||
evictedStr := ""
|
||||
if s.Evicted > 0 {
|
||||
evictedStr = fmt.Sprintf(", %d evicted", s.Evicted)
|
||||
}
|
||||
return fmt.Sprintf("context: ~%d tokens (%d/%d msgs%s)",
|
||||
s.ApproxTokens, s.BoundedMessages, s.StoredMessages, evictedStr)
|
||||
}
|
||||
|
||||
// FormatContextDebug returns a multi-line breakdown of the bounded history
|
||||
// useful for debug output.
|
||||
func (cb *ContextBuilder) FormatContextDebug() string {
|
||||
var sb strings.Builder
|
||||
s := cb.Stats()
|
||||
sb.WriteString(fmt.Sprintf("=== context window: ~%d tokens | %d bounded / %d stored | soft=%d hard=%d ===\n",
|
||||
s.ApproxTokens, s.BoundedMessages, s.StoredMessages, s.SoftLimit/4, s.HardLimit/4))
|
||||
for i, m := range cb.BoundedHistory() {
|
||||
preview := m.Content
|
||||
if len(preview) > 80 {
|
||||
preview = preview[:80] + "..."
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf(" [%d] role=%-10s chars=%-6d %q\n", i, m.Role, msgChars(m), preview))
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
|
@ -4,31 +4,46 @@ import "ollie/backend"
|
|||
|
||||
// Session is an ephemeral in-memory state backend.
|
||||
// It lives only for the duration of the process; nothing is persisted.
|
||||
//
|
||||
// History() returns a bounded context window via ContextBuilder to prevent
|
||||
// prompt explosion across multi-step agent loops. Use NewSessionWithConfig
|
||||
// to override the default limits.
|
||||
type Session struct {
|
||||
goal string
|
||||
history []backend.Message
|
||||
ctx *ContextBuilder
|
||||
complete bool
|
||||
}
|
||||
|
||||
// NewSession creates a Session with the goal already placed as the first
|
||||
// user message in history, so the loop can call backend.Chat immediately.
|
||||
// NewSession creates a Session with default context window limits.
|
||||
func NewSession(goal string) *Session {
|
||||
return &Session{
|
||||
goal: goal,
|
||||
history: []backend.Message{
|
||||
{Role: "user", Content: goal},
|
||||
},
|
||||
}
|
||||
return NewSessionWithConfig(goal, ContextConfig{})
|
||||
}
|
||||
|
||||
func (s *Session) Goal() string { return s.goal }
|
||||
func (s *Session) History() []backend.Message { return s.history }
|
||||
func (s *Session) IsComplete() bool { return s.complete }
|
||||
// NewSessionWithConfig creates a Session with explicit context window limits.
|
||||
// Pass a zero-value ContextConfig to use all defaults.
|
||||
func NewSessionWithConfig(goal string, cfg ContextConfig) *Session {
|
||||
s := &Session{
|
||||
goal: goal,
|
||||
ctx: NewContextBuilder(cfg),
|
||||
}
|
||||
s.ctx.Append(backend.Message{Role: "user", Content: goal})
|
||||
return s
|
||||
}
|
||||
|
||||
func (s *Session) Goal() string { return s.goal }
|
||||
|
||||
// History returns the bounded history window safe for passing to the backend.
|
||||
// A compaction notice is injected when older messages have been evicted.
|
||||
func (s *Session) History() []backend.Message {
|
||||
return s.ctx.BoundedHistoryWithNotice()
|
||||
}
|
||||
|
||||
func (s *Session) IsComplete() bool { return s.complete }
|
||||
|
||||
func (s *Session) Update(assistant backend.Message, results []ToolResult) error {
|
||||
s.history = append(s.history, assistant)
|
||||
s.ctx.Append(assistant)
|
||||
for _, r := range results {
|
||||
s.history = append(s.history, backend.Message{
|
||||
s.ctx.Append(backend.Message{
|
||||
Role: "tool",
|
||||
Content: r.Content,
|
||||
ToolCallID: r.ToolCallID,
|
||||
|
|
@ -46,5 +61,20 @@ func (s *Session) MarkComplete() error {
|
|||
// flag so the loop will run again on the next call to Loop.Run.
|
||||
func (s *Session) AppendUserMessage(content string) {
|
||||
s.complete = false
|
||||
s.history = append(s.history, backend.Message{Role: "user", Content: content})
|
||||
s.ctx.Append(backend.Message{Role: "user", Content: content})
|
||||
}
|
||||
|
||||
// ContextStats returns stats about the current context window.
|
||||
func (s *Session) ContextStats() ContextStats {
|
||||
return s.ctx.Stats()
|
||||
}
|
||||
|
||||
// ContextStatsString returns a one-line human-readable context summary.
|
||||
func (s *Session) ContextStatsString() string {
|
||||
return s.ctx.ContextStatsString()
|
||||
}
|
||||
|
||||
// ContextDebug returns a multi-line breakdown of the bounded history.
|
||||
func (s *Session) ContextDebug() string {
|
||||
return s.ctx.FormatContextDebug()
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue