recursive-llm-ts 4.5.0 → 4.7.0

package/go/rlm/textrank_test.go

@@ -0,0 +1,335 @@
+package rlm
+
+import (
+	"math"
+	"strings"
+	"testing"
+)
+
+// ─── CosineSimilarity Tests ─────────────────────────────────────────────────
+
+func TestCosineSimilarity_IdenticalVectors(t *testing.T) {
+	v := tfidfVector{
+		terms: map[string]float64{"hello": 1.0, "world": 2.0},
+		norm:  math.Sqrt(5.0),
+	}
+	sim := cosineSimilarity(v, v)
+
+	if sim < 0.99 || sim > 1.01 {
+		t.Errorf("expected cosine similarity ~1.0 for identical vectors, got %f", sim)
+	}
+}
+
+func TestCosineSimilarity_OrthogonalVectors(t *testing.T) {
+	a := tfidfVector{
+		terms: map[string]float64{"hello": 1.0},
+		norm:  1.0,
+	}
+	b := tfidfVector{
+		terms: map[string]float64{"world": 1.0},
+		norm:  1.0,
+	}
+	sim := cosineSimilarity(a, b)
+
+	if sim != 0 {
+		t.Errorf("expected cosine similarity 0 for orthogonal vectors, got %f", sim)
+	}
+}
+
+func TestCosineSimilarity_ZeroVector(t *testing.T) {
+	a := tfidfVector{terms: map[string]float64{}, norm: 0}
+	b := tfidfVector{terms: map[string]float64{"hello": 1.0}, norm: 1.0}
+
+	sim := cosineSimilarity(a, b)
+	if sim != 0 {
+		t.Errorf("expected 0 for zero vector, got %f", sim)
+	}
+}
+
+func TestCosineSimilarity_PartialOverlap(t *testing.T) {
+	a := tfidfVector{
+		terms: map[string]float64{"hello": 1.0, "world": 1.0},
+		norm:  math.Sqrt(2.0),
+	}
+	b := tfidfVector{
+		terms: map[string]float64{"hello": 1.0, "foo": 1.0},
+		norm:  math.Sqrt(2.0),
+	}
+	sim := cosineSimilarity(a, b)
+
+	// dot = 1*1 = 1, norm product = sqrt(2)*sqrt(2) = 2, sim = 0.5
+	if sim < 0.49 || sim > 0.51 {
+		t.Errorf("expected cosine similarity ~0.5, got %f", sim)
+	}
+}
+
+// ─── BuildSimilarityGraph Tests ─────────────────────────────────────────────
+
+func TestBuildSimilarityGraph_SimilarSentences(t *testing.T) {
+	sentences := []string{
+		"The machine learning model processes data efficiently.",
+		"The deep learning model processes information quickly.",
+		"The weather today is sunny and warm.",
+	}
+
+	config := DefaultTextRankConfig()
+	config.MinSimilarity = 0.01 // Very low threshold to ensure edges
+	graph := BuildSimilarityGraph(sentences, config)
+
+	if len(graph) != 3 {
+		t.Fatalf("expected 3x3 graph, got %d rows", len(graph))
+	}
+
+	// Sentences 0 and 1 share terms (model, processes, learning) - should have higher similarity
+	// Sentence 2 is about weather - should have lower similarity with 0 and 1
+	sim01 := graph[0][1]
+	sim02 := graph[0][2]
+
+	if sim01 <= 0 {
+		t.Error("expected positive similarity between similar sentences 0 and 1")
+	}
+
+	// 0-1 should be more similar than 0-2
+	if sim01 <= sim02 {
+		t.Errorf("expected sentences 0,1 to be more similar than 0,2: sim01=%f, sim02=%f", sim01, sim02)
+	}
+}
+
+func TestBuildSimilarityGraph_Symmetric(t *testing.T) {
+	sentences := []string{"First sentence here.", "Second sentence there.", "Third different topic."}
+	config := DefaultTextRankConfig()
+	config.MinSimilarity = 0.0
+	graph := BuildSimilarityGraph(sentences, config)
+
+	for i := 0; i < len(graph); i++ {
+		for j := 0; j < len(graph); j++ {
+			if graph[i][j] != graph[j][i] {
+				t.Errorf("graph not symmetric: [%d][%d]=%f != [%d][%d]=%f",
+					i, j, graph[i][j], j, i, graph[j][i])
+			}
+		}
+	}
+}
+
+func TestBuildSimilarityGraph_DiagonalZero(t *testing.T) {
+	sentences := []string{"First.", "Second.", "Third."}
+	config := DefaultTextRankConfig()
+	graph := BuildSimilarityGraph(sentences, config)
+
+	for i := 0; i < len(graph); i++ {
+		if graph[i][i] != 0 {
+			t.Errorf("expected zero self-similarity, got %f for sentence %d", graph[i][i], i)
+		}
+	}
+}
+
+// ─── PageRank Tests ─────────────────────────────────────────────────────────
+
+func TestPageRank_UniformGraph(t *testing.T) {
+	// Fully connected graph with equal weights -> all scores should be equal
+	n := 4
+	graph := make([][]float64, n)
+	for i := range graph {
+		graph[i] = make([]float64, n)
+		for j := range graph[i] {
+			if i != j {
+				graph[i][j] = 1.0
+			}
+		}
+	}
+
+	config := DefaultTextRankConfig()
+	scores := PageRank(graph, config)
+
+	if len(scores) != n {
+		t.Fatalf("expected %d scores, got %d", n, len(scores))
+	}
+
+	// All scores should be approximately equal
+	for i := 1; i < n; i++ {
+		if math.Abs(scores[i]-scores[0]) > 0.01 {
+			t.Errorf("expected uniform scores, got %v", scores)
+			break
+		}
+	}
+}
+
+func TestPageRank_StarGraph(t *testing.T) {
+	// Star graph: node 0 connected to all others, others only connected to 0
+	// Node 0 should have the highest score
+	n := 5
+	graph := make([][]float64, n)
+	for i := range graph {
+		graph[i] = make([]float64, n)
+	}
+	for i := 1; i < n; i++ {
+		graph[0][i] = 1.0
+		graph[i][0] = 1.0
+	}
+
+	config := DefaultTextRankConfig()
+	scores := PageRank(graph, config)
+
+	// Node 0 (hub) should have the highest score
+	for i := 1; i < n; i++ {
+		if scores[0] <= scores[i] {
+			t.Errorf("expected hub node (0) to have highest score: scores[0]=%f, scores[%d]=%f",
+				scores[0], i, scores[i])
+		}
+	}
+}
+
+func TestPageRank_EmptyGraph(t *testing.T) {
+	scores := PageRank(nil, DefaultTextRankConfig())
+	if scores != nil {
+		t.Errorf("expected nil for empty graph, got %v", scores)
+	}
+}
+
+func TestPageRank_DisconnectedGraph(t *testing.T) {
+	// Graph with no edges -> all scores should be equal (from the (1-d)/n term)
+	n := 3
+	graph := make([][]float64, n)
+	for i := range graph {
+		graph[i] = make([]float64, n)
+	}
+
+	config := DefaultTextRankConfig()
+	scores := PageRank(graph, config)
+
+	for i := 1; i < n; i++ {
+		if math.Abs(scores[i]-scores[0]) > 0.001 {
+			t.Errorf("expected equal scores for disconnected graph, got %v", scores)
+			break
+		}
+	}
+}
+
+func TestPageRank_Convergence(t *testing.T) {
+	// PageRank should converge (scores sum to approximately 1)
+	n := 4
+	graph := make([][]float64, n)
+	for i := range graph {
+		graph[i] = make([]float64, n)
+		for j := range graph[i] {
+			if i != j {
+				graph[i][j] = float64((i + j) % 3) // asymmetric weights
+			}
+		}
+	}
+
+	config := DefaultTextRankConfig()
+	scores := PageRank(graph, config)
+
+	sum := 0.0
+	for _, s := range scores {
+		sum += s
+		if s < 0 {
+			t.Errorf("negative PageRank score: %f", s)
+		}
+	}
+
+	// Sum should be approximately 1.0
+	if sum < 0.9 || sum > 1.1 {
+		t.Errorf("expected PageRank scores to sum to ~1.0, got %f", sum)
+	}
+}
+
+// ─── CompressContextTextRank Tests ──────────────────────────────────────────
+
+func TestCompressContextTextRank_NoCompressionNeeded(t *testing.T) {
+	text := "Short text that fits."
+	result := CompressContextTextRank(text, 1000)
+
+	if result != text {
+		t.Errorf("expected unchanged text, got %q", result)
+	}
+}
+
+func TestCompressContextTextRank_CompressesLargeText(t *testing.T) {
+	var parts []string
+	for i := 0; i < 30; i++ {
+		parts = append(parts, "Machine learning algorithms process large datasets to find patterns in the data.")
+	}
+	parts = append(parts, "The stock price of AAPL rose 15% to $198.50 after the earnings report.")
+	parts = append(parts, "Quantum computing achieved 99.9% gate fidelity using topological qubits.")
+	text := strings.Join(parts, " ")
+
+	originalTokens := EstimateTokens(text)
+	targetTokens := originalTokens / 3
+
+	result := CompressContextTextRank(text, targetTokens)
+
+	resultTokens := EstimateTokens(result)
+	if resultTokens > targetTokens+10 {
+		t.Errorf("compressed result (%d tokens) exceeds target (%d tokens)", resultTokens, targetTokens)
+	}
+}
+
+func TestCompressContextTextRank_PreservesOrder(t *testing.T) {
+	text := "Alpha is the first Greek letter. Beta comes after alpha. Gamma is the third letter. Delta is fourth. Epsilon follows delta."
+	result := CompressContextTextRank(text, 30)
+
+	sentences := SplitSentences(result)
+	for i := 1; i < len(sentences); i++ {
+		posI := strings.Index(text, sentences[i])
+		posPrev := strings.Index(text, sentences[i-1])
+		if posI < posPrev {
+			t.Errorf("order not preserved: %q before %q", sentences[i-1], sentences[i])
+		}
+	}
+}
+
+func TestCompressContextTextRank_Empty(t *testing.T) {
+	result := CompressContextTextRank("", 100)
+	if result != "" {
+		t.Errorf("expected empty result, got %q", result)
+	}
+}
+
+func TestCompressContextTextRank_WithConfig(t *testing.T) {
+	text := "Machine learning models. Deep learning networks. Neural network architectures. Data processing pipelines. Cloud computing infrastructure."
+
+	config := TextRankConfig{
+		DampingFactor:        0.85,
+		MaxIterations:        50,
+		ConvergenceThreshold: 0.001,
+		MinSimilarity:        0.0, // Allow all edges
+	}
+
+	result := CompressContextTextRankWithConfig(text, 15, config)
+
+	if len(result) == 0 {
+		t.Error("expected non-empty result")
+	}
+	if EstimateTokens(result) > 20 { // Some slack
+		t.Errorf("result exceeds budget: %d tokens", EstimateTokens(result))
+	}
+}
+
+// ─── TextRank vs TF-IDF Comparison ──────────────────────────────────────────
+
+func TestTextRank_DifferentFromTFIDF(t *testing.T) {
+	// TextRank and TF-IDF should produce different rankings because TextRank
+	// considers inter-sentence relationships while TF-IDF only considers term rarity
+	text := "Machine learning processes data. " +
+		"Deep learning is a subset of machine learning. " +
+		"Neural networks power deep learning. " +
+		"The weather is sunny today. " +
+		"Rain is expected tomorrow."
+
+	tfidfResult := CompressContextTFIDF(text, 20)
+	textrankResult := CompressContextTextRank(text, 20)
+
+	// They should both produce non-empty results
+	if len(tfidfResult) == 0 {
+		t.Error("TF-IDF result is empty")
+	}
+	if len(textrankResult) == 0 {
+		t.Error("TextRank result is empty")
+	}
+
+	// They CAN be the same for simple inputs, but both should work
+	t.Logf("TF-IDF:    %q", tfidfResult)
+	t.Logf("TextRank:  %q", textrankResult)
+}

package/go/rlm/tfidf.go

@@ -0,0 +1,225 @@
+package rlm
+
+import (
+	"math"
+	"sort"
+	"strings"
+	"unicode"
+)
+
+// ─── TF-IDF Extractive Context Compression ──────────────────────────────────
+//
+// Pure Go, zero external dependencies, zero API calls.
+// Scores sentences by Term Frequency - Inverse Document Frequency,
+// selects top-K sentences that fit within a token budget,
+// and preserves original document order.
+
+// ScoredSentence holds a sentence with its TF-IDF score and original position.
+type ScoredSentence struct {
+	Text  string
+	Score float64
+	Index int // original position in the document (for order preservation)
+}
+
+// SplitSentences breaks text into sentences using punctuation boundaries.
+// Handles ". ", "! ", "? " as sentence terminators, plus paragraph breaks.
+func SplitSentences(text string) []string {
+	if len(strings.TrimSpace(text)) == 0 {
+		return nil
+	}
+
+	var sentences []string
+	var current strings.Builder
+
+	runes := []rune(text)
+	for i := 0; i < len(runes); i++ {
+		r := runes[i]
+		current.WriteRune(r)
+
+		// Check for sentence-ending punctuation followed by space/newline/end
+		isSentenceEnd := false
+		if r == '.' || r == '!' || r == '?' {
+			if i+1 >= len(runes) {
+				isSentenceEnd = true
+			} else {
+				next := runes[i+1]
+				if next == ' ' || next == '\n' || next == '\r' || next == '\t' {
+					isSentenceEnd = true
+				}
+			}
+		}
+
+		// Also split on double newline (paragraph boundary)
+		if r == '\n' && i+1 < len(runes) && runes[i+1] == '\n' {
+			s := strings.TrimSpace(current.String())
+			if len(s) > 0 {
+				sentences = append(sentences, s)
+			}
+			current.Reset()
+			i++ // skip the second newline
+			continue
+		}
+
+		if isSentenceEnd {
+			s := strings.TrimSpace(current.String())
+			if len(s) > 0 {
+				sentences = append(sentences, s)
+			}
+			current.Reset()
+		}
+	}
+
+	// Flush remaining text
+	if remaining := strings.TrimSpace(current.String()); len(remaining) > 0 {
+		sentences = append(sentences, remaining)
+	}
+
+	return sentences
+}
+
+// TokenizeWords splits text into lowercase word tokens, filtering non-alphanumeric characters.
+func TokenizeWords(text string) []string {
+	return strings.FieldsFunc(strings.ToLower(text), func(r rune) bool {
+		return !unicode.IsLetter(r) && !unicode.IsDigit(r)
+	})
+}
+
+// stopWords contains common English stop words to optionally filter for better TF-IDF scoring.
+var stopWords = map[string]bool{
+	"a": true, "an": true, "the": true, "and": true, "or": true, "but": true,
+	"in": true, "on": true, "at": true, "to": true, "for": true, "of": true,
+	"with": true, "by": true, "from": true, "is": true, "are": true, "was": true,
+	"were": true, "be": true, "been": true, "being": true, "have": true, "has": true,
+	"had": true, "do": true, "does": true, "did": true, "will": true, "would": true,
+	"could": true, "should": true, "may": true, "might": true, "shall": true,
+	"can": true, "this": true, "that": true, "these": true, "those": true,
+	"it": true, "its": true, "i": true, "you": true, "he": true, "she": true,
+	"we": true, "they": true, "me": true, "him": true, "her": true, "us": true,
+	"them": true, "my": true, "your": true, "his": true, "our": true, "their": true,
+	"not": true, "no": true, "if": true, "as": true, "so": true, "than": true,
+	"then": true, "also": true, "just": true, "about": true, "into": true,
+}
+
+// FilterStopWords removes common stop words from a word list.
+func FilterStopWords(words []string) []string {
+	filtered := make([]string, 0, len(words))
+	for _, w := range words {
+		if !stopWords[w] && len(w) > 1 {
+			filtered = append(filtered, w)
+		}
+	}
+	return filtered
+}
+
+// ComputeTFIDF computes TF-IDF scores for each sentence in a document.
+// Returns ScoredSentence slice with scores and original indices.
+func ComputeTFIDF(sentences []string) []ScoredSentence {
+	if len(sentences) == 0 {
+		return nil
+	}
+
+	// Tokenize each sentence
+	docWords := make([][]string, len(sentences))
+	for i, s := range sentences {
+		words := TokenizeWords(s)
+		docWords[i] = FilterStopWords(words)
+	}
+
+	// Compute document frequency (how many sentences contain each word)
+	df := make(map[string]int)
+	for _, words := range docWords {
+		seen := make(map[string]bool)
+		for _, w := range words {
+			if !seen[w] {
+				df[w]++
+				seen[w] = true
+			}
+		}
+	}
+
+	n := float64(len(sentences))
+
+	// Score each sentence by TF-IDF
+	scored := make([]ScoredSentence, len(sentences))
+	for i, words := range docWords {
+		score := 0.0
+		tf := make(map[string]int)
+		for _, w := range words {
+			tf[w]++
+		}
+		for word, freq := range tf {
+			// TF: term frequency in this sentence
+			// IDF: log(N / df) where N = number of sentences
+			idf := math.Log(n / float64(df[word]))
+			score += float64(freq) * idf
+		}
+		// Normalize by sentence length to avoid bias toward long sentences
+		if len(words) > 0 {
+			score /= math.Sqrt(float64(len(words)))
+		}
+		scored[i] = ScoredSentence{
+			Text:  sentences[i],
+			Score: score,
+			Index: i,
+		}
+	}
+
+	return scored
+}
+
+// CompressContextTFIDF reduces context to fit within a token budget using
+// extractive summarization via TF-IDF sentence scoring.
+// Preserves original sentence order in the output.
+// Returns the original context unchanged if it already fits.
+func CompressContextTFIDF(text string, targetTokens int) string {
+	if EstimateTokens(text) <= targetTokens {
+		return text
+	}
+
+	sentences := SplitSentences(text)
+	if len(sentences) == 0 {
+		return text
+	}
+
+	scored := ComputeTFIDF(sentences)
+
+	// Sort by score descending
+	sort.Slice(scored, func(i, j int) bool {
+		return scored[i].Score > scored[j].Score
+	})
+
+	// Greedily select top sentences until budget is reached
+	var selected []ScoredSentence
+	currentTokens := 0
+	for _, s := range scored {
+		sentTokens := EstimateTokens(s.Text)
+		if currentTokens+sentTokens > targetTokens {
+			continue
+		}
+		selected = append(selected, s)
+		currentTokens += sentTokens
+	}
+
+	if len(selected) == 0 {
+		// Budget too small for even one sentence - truncate the highest-scored
+		if len(scored) > 0 {
+			maxChars := targetTokens * 3 // Conservative chars/token
+			if maxChars > len(scored[0].Text) {
+				maxChars = len(scored[0].Text)
+			}
+			return scored[0].Text[:maxChars]
+		}
+		return text
+	}
+
+	// Re-sort by original index to preserve document order
+	sort.Slice(selected, func(i, j int) bool {
+		return selected[i].Index < selected[j].Index
+	})
+
+	parts := make([]string, len(selected))
+	for i, s := range selected {
+		parts[i] = s.Text
+	}
+	return strings.Join(parts, " ")
+}