recursive-llm-ts 4.8.0 → 5.0.0

package/go/rlm/context_overflow.go

@@ -48,18 +48,18 @@ var modelTokenLimits = map[string]int{
 	"o1-preview":        128000,
 	"o3-mini":           200000,
 	// Anthropic (via LiteLLM/proxy)
-	"claude-3-opus":       200000,
-	"claude-3-sonnet":     200000,
-	"claude-3-haiku":      200000,
-	"claude-3.5-sonnet":   200000,
-	"claude-3.5-haiku":    200000,
-	"claude-sonnet-4":     200000,
-	"claude-opus-4":       200000,
+	"claude-3-opus":     200000,
+	"claude-3-sonnet":   200000,
+	"claude-3-haiku":    200000,
+	"claude-3.5-sonnet": 200000,
+	"claude-3.5-haiku":  200000,
+	"claude-sonnet-4":   200000,
+	"claude-opus-4":     200000,
 	// Llama (common vLLM deployments)
-	"llama-3":     8192,
-	"llama-3.1":   128000,
-	"llama-3.2":   128000,
-	"llama-3.3":   128000,
+	"llama-3":   8192,
+	"llama-3.1": 128000,
+	"llama-3.2": 128000,
+	"llama-3.3": 128000,
 	// Mistral
 	"mistral-7b":    32768,
 	"mixtral-8x7b":  32768,
@@ -181,21 +181,11 @@ func (r *RLM) getResponseTokenBudget() int {
 
 // ─── Token Estimation ────────────────────────────────────────────────────────
 
-// EstimateTokens provides a fast approximation of token count for a string.
-// Uses a character-to-token ratio heuristic. This is intentionally conservative
-// (over-estimates slightly) to avoid overflow.
-//
-// Approximate ratios for common encodings:
-//   - English text: ~4 chars/token (cl100k_base)
-//   - JSON/code:    ~3.5 chars/token
-//   - CJK text:     ~1.5 chars/token
-//   - Mixed:        ~3.5 chars/token (safe default)
+// EstimateTokens returns the token count for a string using the global tokenizer.
+// When SetDefaultTokenizer has been called with a model name, this uses accurate
+// BPE tokenization via tiktoken. Otherwise falls back to a ~3.5 chars/token heuristic.
 func EstimateTokens(text string) int {
-	if len(text) == 0 {
-		return 0
-	}
-	// Use 3.5 chars/token as conservative estimate
-	return (len(text)*10 + 34) / 35 // equivalent to ceil(len/3.5)
+	return GetTokenizer().CountTokens(text)
 }
 
 // EstimateMessagesTokens estimates the total tokens for a set of chat messages.
@@ -431,8 +421,13 @@ func (cr *contextReducer) reduceByMapReduce(query string, chunks []string, model
 			}
 
 			cr.rlm.stats.LlmCalls++
-			summaries[idx] = result
-			cr.obs.Debug("overflow", "Chunk %d/%d summarized: %d -> %d chars", idx+1, len(chunks), len(chunkText), len(result))
+			if result.Usage != nil {
+				cr.rlm.stats.PromptTokens += result.Usage.PromptTokens
+				cr.rlm.stats.CompletionTokens += result.Usage.CompletionTokens
+				cr.rlm.stats.TotalTokens += result.Usage.TotalTokens
+			}
+			summaries[idx] = result.Content
+			cr.obs.Debug("overflow", "Chunk %d/%d summarized: %d -> %d chars", idx+1, len(chunks), len(chunkText), len(result.Content))
 		}(i, chunk)
 	}
 
@@ -472,24 +467,20 @@ func (cr *contextReducer) reduceByMapReduce(query string, chunks []string, model
 }
 
 // reduceByTruncation simply truncates context to fit within the limit.
+// Uses the shared TruncateText utility (compression.go).
 func (cr *contextReducer) reduceByTruncation(context string, modelLimit int, overhead int) (string, error) {
 	cr.obs.Debug("overflow", "Using truncation strategy")
 
 	availableTokens := modelLimit - overhead
-	maxChars := availableTokens * 3 // Conservative chars-to-tokens
+	truncated := TruncateText(context, TruncateTextParams{
+		MaxTokens:  availableTokens,
+		MarkerText: "\n\n[... context truncated due to token limit ...]\n\n",
+	})
 
-	if maxChars >= len(context) {
+	if truncated == context {
 		return context, nil
 	}
 
-	// Keep beginning and end, truncate middle (addresses "lost in the middle" problem)
-	keepFromStart := maxChars * 2 / 3
-	keepFromEnd := maxChars / 3
-
-	truncated := context[:keepFromStart] +
-		"\n\n[... context truncated due to token limit ...]\n\n" +
-		context[len(context)-keepFromEnd:]
-
 	cr.obs.Debug("overflow", "Truncated context: %d -> %d chars", len(context), len(truncated))
 	return truncated, nil
 }
@@ -536,8 +527,13 @@ func (cr *contextReducer) reduceByChunkedExtraction(query string, chunks []strin
 			}
 
 			cr.rlm.stats.LlmCalls++
-			if strings.TrimSpace(result) != "NO_RELEVANT_CONTENT" {
-				results[idx] = result
+			if result.Usage != nil {
+				cr.rlm.stats.PromptTokens += result.Usage.PromptTokens
+				cr.rlm.stats.CompletionTokens += result.Usage.CompletionTokens
+				cr.rlm.stats.TotalTokens += result.Usage.TotalTokens
+			}
+			if strings.TrimSpace(result.Content) != "NO_RELEVANT_CONTENT" {
+				results[idx] = result.Content
 			}
 		}(i, chunk)
 	}
@@ -606,7 +602,7 @@ func (cr *contextReducer) reduceByRefine(query string, chunks []string, modelLim
 		{Role: "user", Content: initialPrompt},
 	}
 
-	currentAnswer, err := CallChatCompletion(ChatRequest{
+	initialResult, err := CallChatCompletion(ChatRequest{
 		Model:       cr.rlm.model,
 		Messages:    messages,
 		APIBase:     cr.rlm.apiBase,
@@ -618,6 +614,12 @@ func (cr *contextReducer) reduceByRefine(query string, chunks []string, modelLim
 		return "", fmt.Errorf("refine initial chunk: %w", err)
 	}
 	cr.rlm.stats.LlmCalls++
+	if initialResult.Usage != nil {
+		cr.rlm.stats.PromptTokens += initialResult.Usage.PromptTokens
+		cr.rlm.stats.CompletionTokens += initialResult.Usage.CompletionTokens
+		cr.rlm.stats.TotalTokens += initialResult.Usage.TotalTokens
+	}
+	currentAnswer := initialResult.Content
 	cr.obs.Debug("overflow", "Refine: initial answer from chunk 1/%d (%d chars)", len(chunks), len(currentAnswer))
 
 	// Phase 2: Refine the answer with each subsequent chunk
@@ -638,7 +640,7 @@ func (cr *contextReducer) reduceByRefine(query string, chunks []string, modelLim
 			{Role: "user", Content: refinePrompt},
 		}
 
-		refined, err := CallChatCompletion(ChatRequest{
+		refineResult, err := CallChatCompletion(ChatRequest{
 			Model:       cr.rlm.model,
 			Messages:    messages,
 			APIBase:     cr.rlm.apiBase,
@@ -652,7 +654,12 @@ func (cr *contextReducer) reduceByRefine(query string, chunks []string, modelLim
 			continue
 		}
 		cr.rlm.stats.LlmCalls++
-		currentAnswer = refined
+		if refineResult.Usage != nil {
+			cr.rlm.stats.PromptTokens += refineResult.Usage.PromptTokens
+			cr.rlm.stats.CompletionTokens += refineResult.Usage.CompletionTokens
+			cr.rlm.stats.TotalTokens += refineResult.Usage.TotalTokens
+		}
+		currentAnswer = refineResult.Content
 		cr.obs.Debug("overflow", "Refine: incorporated chunk %d/%d (%d chars)", i+1, len(chunks), len(currentAnswer))
 	}
 
@@ -700,4 +707,3 @@ func (cr *contextReducer) reduceByTextRank(context string, modelLimit int, overh
 		len(context), len(result), EstimateTokens(context), EstimateTokens(result))
 	return result, nil
 }
-

package/go/rlm/context_savings_test.go

@@ -0,0 +1,387 @@
+package rlm
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+)
+
+func deterministicSentence(idx int) string {
+	topics := []string{
+		"architecture", "testing", "performance", "reliability", "observability",
+		"security", "scalability", "maintainability", "usability", "automation",
+	}
+	details := []string{
+		"input validation", "error handling", "resource limits", "data flow", "boundary conditions",
+		"traceability", "deployment safety", "schema consistency", "latency targets", "integration behavior",
+	}
+
+	topic := topics[idx%len(topics)]
+	detail := details[(idx*7)%len(details)]
+	return fmt.Sprintf("Sentence %d discusses topic %s with details about %s. ", idx, topic, detail)
+}
+
+func generateDeterministicContext(targetTokens int) string {
+	if targetTokens <= 0 {
+		return ""
+	}
+
+	var b strings.Builder
+	total := 0
+	for i := 1; total < targetTokens; i++ {
+		s := deterministicSentence(i)
+		b.WriteString(s)
+		total += EstimateTokens(s)
+	}
+	return b.String()
+}
+
+func fixedEnglishProse500Words() string {
+	words := []string{
+		"software", "teams", "benefit", "from", "clear", "requirements", "because", "stable", "interfaces", "reduce",
+		"rework", "and", "improve", "delivery", "predictability", "when", "engineers", "document", "assumptions", "carefully",
+		"review", "cycles", "become", "faster", "while", "quality", "signals", "remain", "visible", "across",
+		"planning", "implementation", "testing", "and", "maintenance", "phases", "in", "long", "lived", "systems",
+	}
+
+	var b strings.Builder
+	for i := 0; i < 500; i++ {
+		if i > 0 {
+			b.WriteByte(' ')
+		}
+		w := words[i%len(words)]
+		if (i+1)%25 == 0 {
+			w += "."
+		}
+		b.WriteString(w)
+	}
+	return b.String()
+}
+
+func percentDifference(base, compare int) float64 {
+	if base == 0 {
+		return 0
+	}
+	return (float64(compare-base) / float64(base)) * 100
+}
+
+func percentSavings(original, reduced int) float64 {
+	if original <= 0 {
+		return 0
+	}
+	return (float64(original-reduced) / float64(original)) * 100
+}
+
+func yesNo(v bool) string {
+	if v {
+		return "yes"
+	}
+	return "no"
+}
+
+func preservesOriginalSentences(original, reduced string) bool {
+	originalSentences := SplitSentences(original)
+	if len(originalSentences) == 0 {
+		return true
+	}
+
+	origSet := make(map[string]bool, len(originalSentences))
+	for _, s := range originalSentences {
+		origSet[strings.TrimSpace(s)] = true
+	}
+
+	for _, s := range SplitSentences(reduced) {
+		s = strings.TrimSpace(s)
+		if s == "" {
+			continue
+		}
+		if strings.Contains(s, "content truncated") {
+			continue
+		}
+		if !origSet[s] {
+			return false
+		}
+	}
+	return true
+}
+
+func episodeContextCost(episodes []*Episode) int {
+	total := 0
+	for _, ep := range episodes {
+		cost := ep.Tokens
+		if ep.Status != EpisodeActive && ep.SummaryTokens > 0 {
+			cost = ep.SummaryTokens
+		}
+		total += cost
+	}
+	return total
+}
+
+func TestContextSavings_TokenizerAccuracy(t *testing.T) {
+	useHeuristicTokenizerForTest(t)
+
+	bpeTokenizer, err := NewTiktokenTokenizer("gpt-4o")
+	if err != nil {
+		t.Fatalf("failed to create BPE tokenizer: %v", err)
+	}
+
+	goSnippet := `package main
+
+import (
+	"fmt"
+	"strings"
+)
+
+func summarize(items []string) map[string]int {
+	result := map[string]int{}
+	for _, item := range items {
+		normalized := strings.TrimSpace(strings.ToLower(item))
+		if normalized == "" {
+			continue
+		}
+		result[normalized]++
+	}
+	return result
+}
+
+func main() {
+	data := []string{"alpha", "beta", "alpha", "gamma", "beta", "alpha"}
+	stats := summarize(data)
+	fmt.Println("stats:", stats)
+}
+`
+
+	jsonData := `{
+  "project": "recursive-llm-ts",
+  "version": "1.0.0",
+  "features": {
+    "lcm": true,
+    "observability": true,
+    "context_overflow": {
+      "enabled": true,
+      "strategy": "tfidf",
+      "max_reduction_attempts": 3
+    }
+  },
+  "items": [
+    {"id": 1, "name": "alpha", "priority": "high"},
+    {"id": 2, "name": "beta", "priority": "medium"},
+    {"id": 3, "name": "gamma", "priority": "low"}
+  ]
+}`
+
+	cjkText := "这是一个固定的中文测试句子，用于衡量分词稳定性。日本語の固定テスト文を使ってトークン数を比較します。고정된 한국어 문장으로 토큰 계산 결과를 확인합니다。"
+
+	testCases := []struct {
+		name    string
+		content string
+	}{
+		{name: "english_prose", content: fixedEnglishProse500Words()},
+		{name: "go_code", content: goSnippet},
+		{name: "json", content: jsonData},
+		{name: "cjk", content: cjkText},
+	}
+
+	t.Logf("Tokenizer accuracy comparison (heuristic default + direct BPE)")
+	for _, tc := range testCases {
+		heuristic := EstimateTokens(tc.content)
+		bpe := bpeTokenizer.CountTokens(tc.content)
+		chars := len([]rune(tc.content))
+		diffPct := percentDifference(bpe, heuristic)
+
+		t.Logf("type=%-14s chars=%5d heuristic=%5d bpe=%5d diff=%7.2f%%", tc.name, chars, heuristic, bpe, diffPct)
+
+		if heuristic <= 0 {
+			t.Fatalf("heuristic token count should be > 0 for %s", tc.name)
+		}
+		if bpe <= 0 {
+			t.Fatalf("BPE token count should be > 0 for %s", tc.name)
+		}
+	}
+}
+
+func TestContextSavings_FiveLevelEscalation(t *testing.T) {
+	useHeuristicTokenizerForTest(t)
+
+	original := generateDeterministicContext(5000)
+	originalTokens := EstimateTokens(original)
+
+	level3 := CompressContextTFIDF(original, 2000)
+	level4 := CompressContextTextRank(original, 2000)
+	level5 := TruncateText(original, TruncateTextParams{MaxTokens: 2000})
+
+	level3Tokens := EstimateTokens(level3)
+	level4Tokens := EstimateTokens(level4)
+	level5Tokens := EstimateTokens(level5)
+
+	t.Logf("Five-level non-LLM escalation comparison")
+	t.Logf("original_tokens=%d", originalTokens)
+	t.Logf("level=3 strategy=tfidf   tokens=%d reduction=%6.2f%% sentence_preserved=%s", level3Tokens, percentSavings(originalTokens, level3Tokens), yesNo(preservesOriginalSentences(original, level3)))
+	t.Logf("level=4 strategy=textrank tokens=%d reduction=%6.2f%% sentence_preserved=%s", level4Tokens, percentSavings(originalTokens, level4Tokens), yesNo(preservesOriginalSentences(original, level4)))
+	t.Logf("level=5 strategy=truncate tokens=%d reduction=%6.2f%% sentence_preserved=%s", level5Tokens, percentSavings(originalTokens, level5Tokens), yesNo(preservesOriginalSentences(original, level5)))
+
+	if level3Tokens >= originalTokens {
+		t.Fatalf("expected TF-IDF to reduce tokens: original=%d level3=%d", originalTokens, level3Tokens)
+	}
+	if level4Tokens >= originalTokens {
+		t.Fatalf("expected TextRank to reduce tokens: original=%d level4=%d", originalTokens, level4Tokens)
+	}
+	if level5Tokens >= originalTokens {
+		t.Fatalf("expected Truncate to reduce tokens: original=%d level5=%d", originalTokens, level5Tokens)
+	}
+}
+
+func TestContextSavings_EpisodicMemoryBudget(t *testing.T) {
+	useHeuristicTokenizerForTest(t)
+
+	manager := NewEpisodeManager("ctx-savings-episodes", EpisodeConfig{
+		MaxEpisodeMessages:    5,
+		MaxEpisodeTokens:      500,
+		TopicChangeThreshold:  0.5,
+		AutoCompactAfterClose: false,
+	})
+
+	baseTime := time.Date(2024, 1, 2, 3, 4, 5, 0, time.UTC)
+	rawTokens := 0
+	for i := 0; i < 50; i++ {
+		content := fmt.Sprintf("Message %d. %s", i+1, generateDeterministicContext(100))
+		tokens := EstimateTokens(content)
+		rawTokens += tokens
+
+		manager.AddMessage(&StoreMessage{
+			ID:        fmt.Sprintf("msg_%03d", i+1),
+			Role:      RoleUser,
+			Content:   content,
+			Tokens:    tokens,
+			Timestamp: baseTime.Add(time.Duration(i) * time.Minute),
+		})
+	}
+
+	episodes := manager.GetAllEpisodes()
+	t.Logf("episodes_created=%d (expected around 10)", len(episodes))
+	if len(episodes) < 9 || len(episodes) > 11 {
+		t.Fatalf("expected around 10 episodes, got %d", len(episodes))
+	}
+
+	for i := 0; i < len(episodes)-1; i++ {
+		summary := fmt.Sprintf("Episode %d summary. %s", i+1, generateDeterministicContext(30))
+		if err := manager.CompactEpisode(episodes[i].ID, summary); err != nil {
+			t.Fatalf("failed to compact episode %s: %v", episodes[i].ID, err)
+		}
+	}
+
+	budgets := []int{200, 500, 1000, 2000}
+	for _, budget := range budgets {
+		selected := manager.GetEpisodesForContext(budget)
+		contextTokens := episodeContextCost(selected)
+		savings := percentSavings(rawTokens, contextTokens)
+		t.Logf("budget=%4d episodes=%2d context_tokens=%5d raw_tokens=%5d savings=%6.2f%%", budget, len(selected), contextTokens, rawTokens, savings)
+
+		if len(selected) == 0 {
+			t.Fatalf("expected at least one episode for budget %d", budget)
+		}
+	}
+}
+
+func TestContextSavings_AllStrategiesComparison(t *testing.T) {
+	useHeuristicTokenizerForTest(t)
+
+	original := generateDeterministicContext(35000)
+	originalTokens := EstimateTokens(original)
+	target := 16000
+
+	tfidf := CompressContextTFIDF(original, target)
+	textrank := CompressContextTextRank(original, target)
+	truncated := TruncateText(original, TruncateTextParams{MaxTokens: target})
+
+	results := []struct {
+		strategy   string
+		content    string
+		tokens     int
+		preserved  bool
+	}{
+		{strategy: "TF-IDF", content: tfidf, tokens: EstimateTokens(tfidf), preserved: preservesOriginalSentences(original, tfidf)},
+		{strategy: "TextRank", content: textrank, tokens: EstimateTokens(textrank), preserved: preservesOriginalSentences(original, textrank)},
+		{strategy: "Truncate", content: truncated, tokens: EstimateTokens(truncated), preserved: preservesOriginalSentences(original, truncated)},
+	}
+
+	t.Logf("strategy comparison for target=%d tokens (original=%d)", target, originalTokens)
+	t.Logf("strategy   output_tokens  reduction%%  sentence_preserved")
+	for _, r := range results {
+		t.Logf("%-9s %12d  %9.2f%%  %s", r.strategy, r.tokens, percentSavings(originalTokens, r.tokens), yesNo(r.preserved))
+		if r.tokens >= originalTokens {
+			t.Fatalf("strategy %s did not reduce tokens: original=%d output=%d", r.strategy, originalTokens, r.tokens)
+		}
+	}
+}
+
+func TestContextSavings_CombinedPipeline(t *testing.T) {
+	useHeuristicTokenizerForTest(t)
+
+	manager := NewEpisodeManager("ctx-savings-pipeline", EpisodeConfig{
+		MaxEpisodeMessages:    10,
+		MaxEpisodeTokens:      1000000,
+		TopicChangeThreshold:  0.5,
+		AutoCompactAfterClose: false,
+	})
+
+	baseTime := time.Date(2024, 5, 10, 9, 30, 0, 0, time.UTC)
+	messageContentByID := make(map[string]string)
+	rawTokens := 0
+
+	for i := 0; i < 100; i++ {
+		id := fmt.Sprintf("pipeline_msg_%03d", i+1)
+		content := fmt.Sprintf("Message %d segment. %s", i+1, generateDeterministicContext(500))
+		tokens := EstimateTokens(content)
+		rawTokens += tokens
+		messageContentByID[id] = content
+
+		manager.AddMessage(&StoreMessage{
+			ID:        id,
+			Role:      RoleUser,
+			Content:   content,
+			Tokens:    tokens,
+			Timestamp: baseTime.Add(time.Duration(i) * time.Minute),
+		})
+	}
+
+	episodes := manager.GetAllEpisodes()
+	if len(episodes) != 10 {
+		t.Fatalf("expected 10 episodes from 100 messages with MaxEpisodeMessages=10, got %d", len(episodes))
+	}
+
+	afterGrouping := episodeContextCost(episodes)
+
+	for i := 0; i < len(episodes)-1; i++ {
+		ep := episodes[i]
+		var b strings.Builder
+		for _, msgID := range ep.MessageIDs {
+			b.WriteString(messageContentByID[msgID])
+			b.WriteString("\n")
+		}
+		summary := CompressContextTFIDF(b.String(), 300)
+		if err := manager.CompactEpisode(ep.ID, summary); err != nil {
+			t.Fatalf("failed to compact episode %s: %v", ep.ID, err)
+		}
+	}
+
+	afterCompaction := episodeContextCost(manager.GetAllEpisodes())
+	selected := manager.GetEpisodesForContext(8000)
+	afterBudgetSelection := episodeContextCost(selected)
+	totalSavings := percentSavings(rawTokens, afterBudgetSelection)
+
+	t.Logf("Combined pipeline results")
+	t.Logf("original_total_tokens=%d", rawTokens)
+	t.Logf("after_episodic_grouping=%d", afterGrouping)
+	t.Logf("after_compaction=%d", afterCompaction)
+	t.Logf("after_budget_selection=%d", afterBudgetSelection)
+	t.Logf("total_savings=%6.2f%%", totalSavings)
+
+	if afterCompaction >= afterGrouping {
+		t.Fatalf("expected compaction to reduce context tokens: grouped=%d compacted=%d", afterGrouping, afterCompaction)
+	}
+	if afterBudgetSelection > 8000 && len(selected) > 0 && selected[0].Status != EpisodeActive {
+		t.Fatalf("expected selected context <= budget when active episode is not the reason for overflow: selected=%d budget=8000", afterBudgetSelection)
+	}
+}

package/go/rlm/doc.go

@@ -8,13 +8,13 @@
 //
 // To use this package in your Go project:
 //
-//	go get github.com/jbeck018/recursive-llm-ts/go
+//	go get github.com/howlerops/recursive-llm-ts/go
 //
 // # Basic Usage
 //
 // Create an RLM engine and execute a completion:
 //
-//	import "github.com/jbeck018/recursive-llm-ts/go/rlm"
+//	import "github.com/howlerops/recursive-llm-ts/go/rlm"
 //
 //	config := rlm.Config{
 //	    MaxDepth:      5,