feat(memory): add DAG-based context budget compression
Introduce a directed acyclic graph for managing context window budget allocation. Nodes represent message segments with token costs; the compressor selects which nodes to retain or summarize based on importance scores and available token budget.
This commit is contained in:
ZanzyTHEbar
parent
aaaafb723d
commit
b5ab7a6108
4 changed files with 720 additions and 0 deletions
109
pkg/memory/dag/budget.go
Normal file
109
pkg/memory/dag/budget.go
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
package dag
|
||||
|
||||
// BudgetConfig defines the percentage allocation for each context section.
|
||||
// All percentages should sum to 100.
|
||||
type BudgetConfig struct {
|
||||
SystemPromptPct int // % for system prompt (identity, rules, skills)
|
||||
ObservationsPct int // % for observation block
|
||||
KnowledgePct int // % for knowledge block (Focus completions)
|
||||
DAGSummariesPct int // % for DAG compressed history
|
||||
RawTailPct int // % for raw recent messages (uncompressed tail)
|
||||
ToolResultsPct int // % for tool call results
|
||||
}
|
||||
|
||||
// DefaultBudgetConfig returns a balanced allocation.
|
||||
func DefaultBudgetConfig() BudgetConfig {
|
||||
return BudgetConfig{
|
||||
SystemPromptPct: 20,
|
||||
ObservationsPct: 10,
|
||||
KnowledgePct: 5,
|
||||
DAGSummariesPct: 25,
|
||||
RawTailPct: 30,
|
||||
ToolResultsPct: 10,
|
||||
}
|
||||
}
|
||||
|
||||
// Budget represents concrete token allocations computed from config and context window.
|
||||
type Budget struct {
|
||||
Total int
|
||||
SystemPrompt int
|
||||
Observations int
|
||||
Knowledge int
|
||||
DAGSummaries int
|
||||
RawTail int
|
||||
ToolResults int
|
||||
}
|
||||
|
||||
// ComputeBudget calculates token allocations from a context window size and config.
|
||||
func ComputeBudget(contextWindow int, cfg BudgetConfig) Budget {
|
||||
return Budget{
|
||||
Total: contextWindow,
|
||||
SystemPrompt: contextWindow * cfg.SystemPromptPct / 100,
|
||||
Observations: contextWindow * cfg.ObservationsPct / 100,
|
||||
Knowledge: contextWindow * cfg.KnowledgePct / 100,
|
||||
DAGSummaries: contextWindow * cfg.DAGSummariesPct / 100,
|
||||
RawTail: contextWindow * cfg.RawTailPct / 100,
|
||||
ToolResults: contextWindow * cfg.ToolResultsPct / 100,
|
||||
}
|
||||
}
|
||||
|
||||
// Remaining returns tokens available after accounting for used amounts.
|
||||
func (b Budget) Remaining(usedSystem, usedObs, usedKnowledge, usedDAG, usedTail, usedTools int) int {
|
||||
used := usedSystem + usedObs + usedKnowledge + usedDAG + usedTail + usedTools
|
||||
remaining := b.Total - used
|
||||
if remaining < 0 {
|
||||
return 0
|
||||
}
|
||||
return remaining
|
||||
}
|
||||
|
||||
// SelectDAGLevel determines which DAG compression level to use given
|
||||
// the available token budget for DAG summaries.
|
||||
func SelectDAGLevel(d *DAG, budgetTokens int) Level {
|
||||
if d == nil || len(d.Nodes) == 0 {
|
||||
return LevelRaw
|
||||
}
|
||||
|
||||
sessionTokens := d.TotalTokens(LevelSession)
|
||||
if sessionTokens > 0 && sessionTokens <= budgetTokens {
|
||||
sectionTokens := d.TotalTokens(LevelSection)
|
||||
if sectionTokens > 0 && sectionTokens <= budgetTokens {
|
||||
chunkTokens := d.TotalTokens(LevelChunk)
|
||||
if chunkTokens <= budgetTokens {
|
||||
return LevelChunk
|
||||
}
|
||||
return LevelSection
|
||||
}
|
||||
return LevelSession
|
||||
}
|
||||
|
||||
return LevelSession
|
||||
}
|
||||
|
||||
// RenderDAGForBudget renders DAG nodes at the most detailed level
|
||||
// that fits within the given token budget.
|
||||
func RenderDAGForBudget(d *DAG, budgetTokens int) string {
|
||||
if d == nil || len(d.Nodes) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
level := SelectDAGLevel(d, budgetTokens)
|
||||
return d.FormatLevel(level)
|
||||
}
|
||||
|
||||
// TailMessageCount estimates how many raw messages fit in the tail budget.
|
||||
// Uses a rough average of ~50 tokens per message.
|
||||
func TailMessageCount(tailBudget int) int {
|
||||
const (
|
||||
avgTokensPerMessage = 50
|
||||
minTail = 4
|
||||
)
|
||||
if tailBudget <= 0 {
|
||||
return minTail
|
||||
}
|
||||
count := tailBudget / avgTokensPerMessage
|
||||
if count < minTail {
|
||||
return minTail
|
||||
}
|
||||
return count
|
||||
}
|
||||
232
pkg/memory/dag/compress.go
Normal file
232
pkg/memory/dag/compress.go
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
package dag
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Message is a minimal message representation for DAG compression.
|
||||
type Message struct {
|
||||
Role string
|
||||
Content string
|
||||
}
|
||||
|
||||
// CompressorConfig controls the deterministic compression behavior.
|
||||
type CompressorConfig struct {
|
||||
ChunkSize int // Messages per chunk node (default 8)
|
||||
SectionSize int // Chunks per section node (default 4)
|
||||
MaxSentences int // Max sentences to extract per message (default 2)
|
||||
TargetRatio float64 // Target compression ratio (default 0.25 = 4:1)
|
||||
}
|
||||
|
||||
// DefaultCompressorConfig returns sensible defaults.
|
||||
func DefaultCompressorConfig() CompressorConfig {
|
||||
return CompressorConfig{
|
||||
ChunkSize: 8,
|
||||
SectionSize: 4,
|
||||
MaxSentences: 2,
|
||||
TargetRatio: 0.25,
|
||||
}
|
||||
}
|
||||
|
||||
// Compressor builds a DAG from raw messages using deterministic
|
||||
// extractive summarization. No LLM calls — fully reproducible.
|
||||
type Compressor struct {
|
||||
cfg CompressorConfig
|
||||
counter int
|
||||
}
|
||||
|
||||
// NewCompressor creates a new deterministic compressor.
|
||||
func NewCompressor(cfg CompressorConfig) *Compressor {
|
||||
return &Compressor{cfg: cfg}
|
||||
}
|
||||
|
||||
// Compress builds a hierarchical DAG from the given messages.
|
||||
// Messages are grouped into chunks, chunks into sections, and
|
||||
// sections into a session summary.
|
||||
func (c *Compressor) Compress(msgs []Message) *DAG {
|
||||
d := NewDAG()
|
||||
if len(msgs) == 0 {
|
||||
return d
|
||||
}
|
||||
|
||||
// Level 1: Chunk summaries
|
||||
chunkNodes := c.buildChunks(msgs, d)
|
||||
if len(chunkNodes) == 0 {
|
||||
return d
|
||||
}
|
||||
|
||||
// Level 2: Section summaries (groups of chunks)
|
||||
sectionNodes := c.buildSections(chunkNodes, d)
|
||||
|
||||
// Level 3: Session summary (if multiple sections)
|
||||
if len(sectionNodes) > 1 {
|
||||
sessionNode := c.buildSessionSummary(sectionNodes, d)
|
||||
d.SetRoots([]string{sessionNode.ID})
|
||||
} else if len(sectionNodes) == 1 {
|
||||
d.SetRoots([]string{sectionNodes[0].ID})
|
||||
} else {
|
||||
ids := make([]string, len(chunkNodes))
|
||||
for i, n := range chunkNodes {
|
||||
ids[i] = n.ID
|
||||
}
|
||||
d.SetRoots(ids)
|
||||
}
|
||||
|
||||
return d
|
||||
}
|
||||
|
||||
func (c *Compressor) nextID(prefix string) string {
|
||||
c.counter++
|
||||
return fmt.Sprintf("%s-%d", prefix, c.counter)
|
||||
}
|
||||
|
||||
func (c *Compressor) buildChunks(msgs []Message, d *DAG) []*Node {
|
||||
var chunks []*Node
|
||||
for i := 0; i < len(msgs); i += c.cfg.ChunkSize {
|
||||
end := i + c.cfg.ChunkSize
|
||||
if end > len(msgs) {
|
||||
end = len(msgs)
|
||||
}
|
||||
|
||||
chunk := msgs[i:end]
|
||||
summary := c.extractChunkSummary(chunk)
|
||||
node := &Node{
|
||||
ID: c.nextID("chunk"),
|
||||
Level: LevelChunk,
|
||||
Summary: summary,
|
||||
Tokens: estimateTokens(summary),
|
||||
StartIdx: i,
|
||||
EndIdx: end,
|
||||
}
|
||||
d.Add(node)
|
||||
chunks = append(chunks, node)
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
func (c *Compressor) buildSections(chunks []*Node, d *DAG) []*Node {
|
||||
if len(chunks) <= c.cfg.SectionSize {
|
||||
return chunks
|
||||
}
|
||||
|
||||
var sections []*Node
|
||||
for i := 0; i < len(chunks); i += c.cfg.SectionSize {
|
||||
end := i + c.cfg.SectionSize
|
||||
if end > len(chunks) {
|
||||
end = len(chunks)
|
||||
}
|
||||
|
||||
group := chunks[i:end]
|
||||
childIDs := make([]string, len(group))
|
||||
var summaryParts []string
|
||||
for j, ch := range group {
|
||||
childIDs[j] = ch.ID
|
||||
summaryParts = append(summaryParts, ch.Summary)
|
||||
}
|
||||
|
||||
combined := strings.Join(summaryParts, " ")
|
||||
summary := extractSentences(combined, c.cfg.MaxSentences)
|
||||
node := &Node{
|
||||
ID: c.nextID("section"),
|
||||
Level: LevelSection,
|
||||
Summary: summary,
|
||||
Tokens: estimateTokens(summary),
|
||||
StartIdx: group[0].StartIdx,
|
||||
EndIdx: group[len(group)-1].EndIdx,
|
||||
Children: childIDs,
|
||||
}
|
||||
d.Add(node)
|
||||
sections = append(sections, node)
|
||||
}
|
||||
return sections
|
||||
}
|
||||
|
||||
func (c *Compressor) buildSessionSummary(sections []*Node, d *DAG) *Node {
|
||||
childIDs := make([]string, len(sections))
|
||||
var summaryParts []string
|
||||
for i, s := range sections {
|
||||
childIDs[i] = s.ID
|
||||
summaryParts = append(summaryParts, s.Summary)
|
||||
}
|
||||
|
||||
combined := strings.Join(summaryParts, " ")
|
||||
summary := extractSentences(combined, c.cfg.MaxSentences)
|
||||
node := &Node{
|
||||
ID: c.nextID("session"),
|
||||
Level: LevelSession,
|
||||
Summary: summary,
|
||||
Tokens: estimateTokens(summary),
|
||||
StartIdx: sections[0].StartIdx,
|
||||
EndIdx: sections[len(sections)-1].EndIdx,
|
||||
Children: childIDs,
|
||||
}
|
||||
d.Add(node)
|
||||
return node
|
||||
}
|
||||
|
||||
// extractChunkSummary produces a deterministic summary of a message chunk.
|
||||
// Strategy: for each message, take the first sentence (or first N chars if short).
|
||||
// Prefix with role to preserve conversational structure.
|
||||
func (c *Compressor) extractChunkSummary(msgs []Message) string {
|
||||
var parts []string
|
||||
for _, m := range msgs {
|
||||
sentence := extractSentences(m.Content, 1)
|
||||
if sentence == "" {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, fmt.Sprintf("%s: %s", m.Role, sentence))
|
||||
}
|
||||
return strings.Join(parts, " | ")
|
||||
}
|
||||
|
||||
// extractSentences pulls the first N sentences from text.
|
||||
// A sentence ends at '.', '!', '?', or '\n\n'.
|
||||
func extractSentences(text string, n int) string {
|
||||
if n <= 0 || text == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
text = strings.TrimSpace(text)
|
||||
var result []string
|
||||
remaining := text
|
||||
|
||||
for i := 0; i < n && remaining != ""; i++ {
|
||||
idx := findSentenceEnd(remaining)
|
||||
if idx < 0 {
|
||||
result = append(result, strings.TrimSpace(remaining))
|
||||
break
|
||||
}
|
||||
sentence := strings.TrimSpace(remaining[:idx+1])
|
||||
if sentence != "" {
|
||||
result = append(result, sentence)
|
||||
}
|
||||
remaining = strings.TrimSpace(remaining[idx+1:])
|
||||
}
|
||||
|
||||
joined := strings.Join(result, " ")
|
||||
const maxLen = 200
|
||||
if utf8.RuneCountInString(joined) > maxLen {
|
||||
runes := []rune(joined)
|
||||
return string(runes[:maxLen]) + "..."
|
||||
}
|
||||
return joined
|
||||
}
|
||||
|
||||
func findSentenceEnd(s string) int {
|
||||
for i, r := range s {
|
||||
if r == '.' || r == '!' || r == '?' {
|
||||
return i
|
||||
}
|
||||
if r == '\n' && i+1 < len(s) && s[i+1] == '\n' {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// estimateTokens provides a rough token count (chars * 2/5 heuristic).
|
||||
func estimateTokens(s string) int {
|
||||
return utf8.RuneCountInString(s) * 2 / 5
|
||||
}
|
||||
245
pkg/memory/dag/dag_test.go
Normal file
245
pkg/memory/dag/dag_test.go
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
package dag
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func makeMessages(n int) []Message {
|
||||
msgs := make([]Message, n)
|
||||
for i := range msgs {
|
||||
role := "user"
|
||||
if i%2 == 1 {
|
||||
role = "assistant"
|
||||
}
|
||||
msgs[i] = Message{
|
||||
Role: role,
|
||||
Content: strings.Repeat("word ", 20) + ".",
|
||||
}
|
||||
}
|
||||
return msgs
|
||||
}
|
||||
|
||||
func TestCompressor_EmptyInput(t *testing.T) {
|
||||
c := NewCompressor(DefaultCompressorConfig())
|
||||
d := c.Compress(nil)
|
||||
assert.Empty(t, d.Nodes)
|
||||
assert.Empty(t, d.Roots)
|
||||
}
|
||||
|
||||
func TestCompressor_SmallInput(t *testing.T) {
|
||||
c := NewCompressor(DefaultCompressorConfig())
|
||||
msgs := []Message{
|
||||
{Role: "user", Content: "Hello, how are you?"},
|
||||
{Role: "assistant", Content: "I'm doing well. Thanks for asking!"},
|
||||
}
|
||||
d := c.Compress(msgs)
|
||||
|
||||
require.Len(t, d.Nodes, 1)
|
||||
chunk := d.NodesAtLevel(LevelChunk)
|
||||
require.Len(t, chunk, 1)
|
||||
assert.Equal(t, 0, chunk[0].StartIdx)
|
||||
assert.Equal(t, 2, chunk[0].EndIdx)
|
||||
assert.Contains(t, chunk[0].Summary, "user:")
|
||||
assert.Contains(t, chunk[0].Summary, "assistant:")
|
||||
}
|
||||
|
||||
func TestCompressor_ChunkSplitting(t *testing.T) {
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
msgs := makeMessages(12)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
chunks := d.NodesAtLevel(LevelChunk)
|
||||
require.Len(t, chunks, 3)
|
||||
|
||||
assert.Equal(t, 0, chunks[0].StartIdx)
|
||||
assert.Equal(t, 4, chunks[0].EndIdx)
|
||||
assert.Equal(t, 4, chunks[1].StartIdx)
|
||||
assert.Equal(t, 8, chunks[1].EndIdx)
|
||||
assert.Equal(t, 8, chunks[2].StartIdx)
|
||||
assert.Equal(t, 12, chunks[2].EndIdx)
|
||||
}
|
||||
|
||||
func TestCompressor_SectionBuilding(t *testing.T) {
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
cfg.SectionSize = 2
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
// 20 messages = 5 chunks, section_size=2 => 3 sections
|
||||
msgs := makeMessages(20)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
chunks := d.NodesAtLevel(LevelChunk)
|
||||
assert.Len(t, chunks, 5)
|
||||
|
||||
sections := d.NodesAtLevel(LevelSection)
|
||||
assert.Len(t, sections, 3)
|
||||
|
||||
// First section covers chunks 0-1 (msgs 0-7)
|
||||
assert.Equal(t, 0, sections[0].StartIdx)
|
||||
assert.Equal(t, 8, sections[0].EndIdx)
|
||||
assert.Len(t, sections[0].Children, 2)
|
||||
}
|
||||
|
||||
func TestCompressor_SessionSummary(t *testing.T) {
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
cfg.SectionSize = 2
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
// 24 messages => 6 chunks => 3 sections => 1 session
|
||||
msgs := makeMessages(24)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
sessions := d.NodesAtLevel(LevelSession)
|
||||
require.Len(t, sessions, 1)
|
||||
assert.Equal(t, 0, sessions[0].StartIdx)
|
||||
assert.Equal(t, 24, sessions[0].EndIdx)
|
||||
assert.Len(t, sessions[0].Children, 3)
|
||||
|
||||
require.Len(t, d.Roots, 1)
|
||||
assert.Equal(t, sessions[0].ID, d.Roots[0])
|
||||
}
|
||||
|
||||
func TestExtractSentences(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
text string
|
||||
n int
|
||||
want string
|
||||
}{
|
||||
{"single sentence", "Hello world.", 1, "Hello world."},
|
||||
{"two sentences", "First sentence. Second sentence.", 2, "First sentence. Second sentence."},
|
||||
{"extract one from many", "A. B. C. D.", 1, "A."},
|
||||
{"empty", "", 1, ""},
|
||||
{"no period", "Hello world", 1, "Hello world"},
|
||||
{"question mark", "What? Who knows.", 1, "What?"},
|
||||
{"exclamation", "Wow! Amazing.", 2, "Wow! Amazing."},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := extractSentences(tt.text, tt.n)
|
||||
assert.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractSentences_Truncation(t *testing.T) {
|
||||
long := strings.Repeat("This is a very long sentence with many words. ", 20)
|
||||
result := extractSentences(long, 5)
|
||||
assert.LessOrEqual(t, len([]rune(result)), 210)
|
||||
assert.True(t, strings.HasSuffix(result, "..."))
|
||||
}
|
||||
|
||||
func TestNode_FormatForPrompt(t *testing.T) {
|
||||
n := &Node{
|
||||
ID: "chunk-1",
|
||||
Level: LevelChunk,
|
||||
Summary: "User asked about auth. Assistant explained JWT flow.",
|
||||
StartIdx: 0,
|
||||
EndIdx: 8,
|
||||
}
|
||||
formatted := n.FormatForPrompt()
|
||||
assert.Contains(t, formatted, "[chunk msgs 0-7]")
|
||||
assert.Contains(t, formatted, "User asked about auth")
|
||||
}
|
||||
|
||||
func TestDAG_FormatLevel(t *testing.T) {
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
msgs := makeMessages(8)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
output := d.FormatLevel(LevelChunk)
|
||||
assert.NotEmpty(t, output)
|
||||
lines := strings.Split(strings.TrimSpace(output), "\n")
|
||||
assert.Len(t, lines, 2)
|
||||
}
|
||||
|
||||
func TestDAG_TotalTokens(t *testing.T) {
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
msgs := makeMessages(12)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
chunkTokens := d.TotalTokens(LevelChunk)
|
||||
assert.Greater(t, chunkTokens, 0)
|
||||
}
|
||||
|
||||
func TestComputeBudget(t *testing.T) {
|
||||
cfg := DefaultBudgetConfig()
|
||||
b := ComputeBudget(100000, cfg)
|
||||
|
||||
assert.Equal(t, 100000, b.Total)
|
||||
assert.Equal(t, 20000, b.SystemPrompt)
|
||||
assert.Equal(t, 10000, b.Observations)
|
||||
assert.Equal(t, 5000, b.Knowledge)
|
||||
assert.Equal(t, 25000, b.DAGSummaries)
|
||||
assert.Equal(t, 30000, b.RawTail)
|
||||
assert.Equal(t, 10000, b.ToolResults)
|
||||
}
|
||||
|
||||
func TestBudget_Remaining(t *testing.T) {
|
||||
b := Budget{Total: 10000}
|
||||
assert.Equal(t, 7000, b.Remaining(1000, 500, 500, 500, 500, 0))
|
||||
assert.Equal(t, 0, b.Remaining(5000, 3000, 1000, 1000, 1000, 0))
|
||||
}
|
||||
|
||||
func TestSelectDAGLevel(t *testing.T) {
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
cfg.SectionSize = 2
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
msgs := makeMessages(24)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
chunkTokens := d.TotalTokens(LevelChunk)
|
||||
|
||||
// Large budget -> most detailed (chunk)
|
||||
assert.Equal(t, LevelChunk, SelectDAGLevel(d, chunkTokens+1000))
|
||||
|
||||
// Very small budget -> session level
|
||||
assert.Equal(t, LevelSession, SelectDAGLevel(d, 10))
|
||||
|
||||
// Nil DAG
|
||||
assert.Equal(t, LevelRaw, SelectDAGLevel(nil, 1000))
|
||||
}
|
||||
|
||||
func TestTailMessageCount(t *testing.T) {
|
||||
assert.Equal(t, 4, TailMessageCount(100)) // Minimum
|
||||
assert.Equal(t, 20, TailMessageCount(1000)) // 1000/50
|
||||
assert.Equal(t, 4, TailMessageCount(0)) // Zero budget
|
||||
assert.Equal(t, 4, TailMessageCount(-1)) // Negative
|
||||
}
|
||||
|
||||
func TestRenderDAGForBudget(t *testing.T) {
|
||||
assert.Empty(t, RenderDAGForBudget(nil, 1000))
|
||||
|
||||
cfg := DefaultCompressorConfig()
|
||||
cfg.ChunkSize = 4
|
||||
cfg.SectionSize = 2
|
||||
c := NewCompressor(cfg)
|
||||
|
||||
// 16 messages => 4 chunks => 2 sections => 1 session
|
||||
msgs := makeMessages(16)
|
||||
d := c.Compress(msgs)
|
||||
|
||||
// Large budget should get chunk-level detail
|
||||
chunkTokens := d.TotalTokens(LevelChunk)
|
||||
output := RenderDAGForBudget(d, chunkTokens+1000)
|
||||
assert.NotEmpty(t, output)
|
||||
assert.Contains(t, output, "[chunk")
|
||||
}
|
||||
134
pkg/memory/dag/node.go
Normal file
134
pkg/memory/dag/node.go
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
package dag
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Level represents the compression tier of a DAG node.
|
||||
type Level int
|
||||
|
||||
const (
|
||||
LevelRaw Level = 0 // Original messages (not stored as nodes)
|
||||
LevelChunk Level = 1 // Chunk summary (~8-16 messages)
|
||||
LevelSection Level = 2 // Section summary (group of chunks)
|
||||
LevelSession Level = 3 // Session summary (top-level)
|
||||
)
|
||||
|
||||
func (l Level) String() string {
|
||||
switch l {
|
||||
case LevelRaw:
|
||||
return "raw"
|
||||
case LevelChunk:
|
||||
return "chunk"
|
||||
case LevelSection:
|
||||
return "section"
|
||||
case LevelSession:
|
||||
return "session"
|
||||
default:
|
||||
return fmt.Sprintf("level-%d", l)
|
||||
}
|
||||
}
|
||||
|
||||
// Node is a single node in the compression DAG. Each node stores an
|
||||
// extractive summary and retains lossless pointers back to the
|
||||
// original message range it covers.
|
||||
type Node struct {
|
||||
ID string `json:"id"`
|
||||
Level Level `json:"level"`
|
||||
Summary string `json:"summary"`
|
||||
Tokens int `json:"tokens"`
|
||||
StartIdx int `json:"start_idx"` // Inclusive index into original message slice
|
||||
EndIdx int `json:"end_idx"` // Exclusive index into original message slice
|
||||
Children []string `json:"children,omitempty"` // Child node IDs (lower level)
|
||||
}
|
||||
|
||||
// MessageRange returns the [start, end) range of original messages this node covers.
|
||||
func (n *Node) MessageRange() (int, int) {
|
||||
return n.StartIdx, n.EndIdx
|
||||
}
|
||||
|
||||
// Span returns how many original messages this node covers.
|
||||
func (n *Node) Span() int {
|
||||
return n.EndIdx - n.StartIdx
|
||||
}
|
||||
|
||||
// FormatForPrompt renders the node as a compact block for context injection.
|
||||
func (n *Node) FormatForPrompt() string {
|
||||
return fmt.Sprintf("[%s msgs %d-%d] %s", n.Level, n.StartIdx, n.EndIdx-1, n.Summary)
|
||||
}
|
||||
|
||||
// DAG is the hierarchical compression tree. Nodes at higher levels
|
||||
// summarize groups of lower-level nodes. The root level covers the
|
||||
// entire session.
|
||||
type DAG struct {
|
||||
Nodes map[string]*Node `json:"nodes"`
|
||||
Roots []string `json:"roots"` // Top-level node IDs (highest compression)
|
||||
}
|
||||
|
||||
// NewDAG creates an empty DAG.
|
||||
func NewDAG() *DAG {
|
||||
return &DAG{
|
||||
Nodes: make(map[string]*Node),
|
||||
}
|
||||
}
|
||||
|
||||
// Add inserts a node into the DAG.
|
||||
func (d *DAG) Add(node *Node) {
|
||||
d.Nodes[node.ID] = node
|
||||
}
|
||||
|
||||
// SetRoots sets the top-level node IDs.
|
||||
func (d *DAG) SetRoots(ids []string) {
|
||||
d.Roots = ids
|
||||
}
|
||||
|
||||
// Get returns a node by ID, or nil if not found.
|
||||
func (d *DAG) Get(id string) *Node {
|
||||
return d.Nodes[id]
|
||||
}
|
||||
|
||||
// NodesAtLevel returns all nodes at the given compression level, ordered by StartIdx.
|
||||
func (d *DAG) NodesAtLevel(level Level) []*Node {
|
||||
var result []*Node
|
||||
for _, n := range d.Nodes {
|
||||
if n.Level == level {
|
||||
result = append(result, n)
|
||||
}
|
||||
}
|
||||
sortByStart(result)
|
||||
return result
|
||||
}
|
||||
|
||||
// TotalTokens returns the sum of tokens across all nodes at the given level.
|
||||
func (d *DAG) TotalTokens(level Level) int {
|
||||
total := 0
|
||||
for _, n := range d.Nodes {
|
||||
if n.Level == level {
|
||||
total += n.Tokens
|
||||
}
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// FormatLevel renders all nodes at a given level as a prompt-ready string.
|
||||
func (d *DAG) FormatLevel(level Level) string {
|
||||
nodes := d.NodesAtLevel(level)
|
||||
if len(nodes) == 0 {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
for _, n := range nodes {
|
||||
sb.WriteString(n.FormatForPrompt())
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func sortByStart(nodes []*Node) {
|
||||
for i := 1; i < len(nodes); i++ {
|
||||
for j := i; j > 0 && nodes[j].StartIdx < nodes[j-1].StartIdx; j-- {
|
||||
nodes[j], nodes[j-1] = nodes[j-1], nodes[j]
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue