recursive-llm-ts 4.9.0 → 5.0.0

package/go/rlm/lcm_agentic_map.go

@@ -0,0 +1,317 @@
+package rlm
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// ─── Agentic-Map Operator ───────────────────────────────────────────────────
+// Implements Operator-Level Recursion from the LCM paper (Section 3.1).
+// Similar to LLM-Map, but spawns a full sub-agent session for each item
+// with access to tools (file I/O, code execution, multi-step reasoning).
+
+// AgenticMapConfig configures an Agentic-Map operation.
+type AgenticMapConfig struct {
+	InputPath    string      `json:"input_path"`    // Path to JSONL input file
+	OutputPath   string      `json:"output_path"`   // Path to JSONL output file
+	Prompt       string      `json:"prompt"`        // Prompt template ({{item}} is replaced)
+	OutputSchema *JSONSchema `json:"output_schema"` // Schema for output validation
+	Concurrency  int         `json:"concurrency"`   // Worker pool size (default: 8, lower than LLM-Map)
+	MaxRetries   int         `json:"max_retries"`   // Per-item retry limit (default: 2)
+	Model        string      `json:"model"`         // Model for sub-agents (default: engine model)
+	ReadOnly     bool        `json:"read_only"`     // If true, sub-agents cannot modify filesystem
+	MaxDepth     int         `json:"max_depth"`     // Max recursion depth for sub-agents (default: 3)
+	MaxIter      int         `json:"max_iterations"` // Max iterations per sub-agent (default: 15)
+}
+
+// AgenticMapResult contains results of an Agentic-Map operation.
+type AgenticMapResult struct {
+	TotalItems  int                 `json:"total_items"`
+	Completed   int                 `json:"completed"`
+	Failed      int                 `json:"failed"`
+	OutputPath  string              `json:"output_path"`
+	Duration    time.Duration       `json:"duration"`
+	TokensUsed  int                 `json:"tokens_used"`
+	ItemResults []AgenticItemResult `json:"item_results,omitempty"`
+}
+
+// AgenticItemResult tracks the status of a single agentic-map item.
+type AgenticItemResult struct {
+	Index      int             `json:"index"`
+	Status     MapItemStatus   `json:"status"`
+	Output     json.RawMessage `json:"output,omitempty"`
+	Error      string          `json:"error,omitempty"`
+	Retries    int             `json:"retries"`
+	LLMCalls   int             `json:"llm_calls"`
+	Iterations int             `json:"iterations"`
+}
+
+// AgenticMapper executes Agentic-Map operations using full sub-agent sessions.
+type AgenticMapper struct {
+	model       string
+	apiBase     string
+	apiKey      string
+	timeout     int
+	extraParams map[string]interface{}
+	observer    *Observer
+}
+
+// NewAgenticMapper creates a new Agentic-Map executor.
+func NewAgenticMapper(model, apiBase, apiKey string, timeout int, extraParams map[string]interface{}, observer *Observer) *AgenticMapper {
+	return &AgenticMapper{
+		model:       model,
+		apiBase:     apiBase,
+		apiKey:      apiKey,
+		timeout:     timeout,
+		extraParams: extraParams,
+		observer:    observer,
+	}
+}
+
+// Execute runs an Agentic-Map operation: parallel sub-agent sessions over JSONL input.
+func (am *AgenticMapper) Execute(config AgenticMapConfig) (*AgenticMapResult, error) {
+	start := time.Now()
+
+	// Apply defaults
+	if config.Concurrency <= 0 {
+		config.Concurrency = 8 // Lower default than LLM-Map due to heavier per-item cost
+	}
+	if config.MaxRetries <= 0 {
+		config.MaxRetries = 2
+	}
+	if config.MaxDepth <= 0 {
+		config.MaxDepth = 3
+	}
+	if config.MaxIter <= 0 {
+		config.MaxIter = 15
+	}
+	model := config.Model
+	if model == "" {
+		model = am.model
+	}
+
+	am.observer.Debug("lcm.agentic_map", "Starting Agentic-Map: input=%s, concurrency=%d, model=%s, read_only=%v",
+		config.InputPath, config.Concurrency, model, config.ReadOnly)
+
+	// Read input items
+	items, err := readJSONLFile(config.InputPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read input: %w", err)
+	}
+
+	am.observer.Debug("lcm.agentic_map", "Read %d items from %s", len(items), config.InputPath)
+
+	// Initialize results
+	results := make([]AgenticItemResult, len(items))
+	for i := range results {
+		results[i] = AgenticItemResult{
+			Index:  i,
+			Status: MapItemPending,
+		}
+	}
+
+	// Worker pool
+	var wg sync.WaitGroup
+	itemChan := make(chan int, len(items))
+	var totalTokens int64
+
+	for i := range items {
+		itemChan <- i
+	}
+	close(itemChan)
+
+	var mu sync.Mutex
+	for w := 0; w < config.Concurrency && w < len(items); w++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for idx := range itemChan {
+				result := am.processAgenticItem(items[idx], config, model)
+				atomic.AddInt64(&totalTokens, int64(result.tokensUsed))
+
+				mu.Lock()
+				results[idx] = AgenticItemResult{
+					Index:      idx,
+					Status:     result.status,
+					Output:     result.output,
+					Error:      result.errMsg,
+					Retries:    result.retries,
+					LLMCalls:   result.llmCalls,
+					Iterations: result.iterations,
+				}
+				mu.Unlock()
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	// Write output file
+	if config.OutputPath != "" {
+		if err := writeAgenticOutput(config.OutputPath, results); err != nil {
+			return nil, fmt.Errorf("failed to write output: %w", err)
+		}
+	}
+
+	completed, failed := 0, 0
+	for _, r := range results {
+		switch r.Status {
+		case MapItemCompleted:
+			completed++
+		case MapItemFailed:
+			failed++
+		}
+	}
+
+	duration := time.Since(start)
+	am.observer.Debug("lcm.agentic_map", "Agentic-Map complete: %d/%d succeeded, %d failed in %s",
+		completed, len(items), failed, duration)
+	am.observer.Event("lcm.agentic_map_complete", map[string]string{
+		"total_items": fmt.Sprintf("%d", len(items)),
+		"completed":   fmt.Sprintf("%d", completed),
+		"failed":      fmt.Sprintf("%d", failed),
+		"duration_ms": fmt.Sprintf("%d", duration.Milliseconds()),
+		"tokens_used": fmt.Sprintf("%d", totalTokens),
+	})
+
+	return &AgenticMapResult{
+		TotalItems:  len(items),
+		Completed:   completed,
+		Failed:      failed,
+		OutputPath:  config.OutputPath,
+		Duration:    duration,
+		TokensUsed:  int(totalTokens),
+		ItemResults: results,
+	}, nil
+}
+
+// ─── Per-Item Sub-Agent Processing ──────────────────────────────────────────
+
+type agenticItemResult struct {
+	status     MapItemStatus
+	output     json.RawMessage
+	errMsg     string
+	retries    int
+	tokensUsed int
+	llmCalls   int
+	iterations int
+}
+
+func (am *AgenticMapper) processAgenticItem(item json.RawMessage, config AgenticMapConfig, model string) agenticItemResult {
+	prompt := strings.ReplaceAll(config.Prompt, "{{item}}", string(item))
+
+	var lastErr string
+	for attempt := 0; attempt <= config.MaxRetries; attempt++ {
+		currentPrompt := prompt
+		if attempt > 0 && lastErr != "" {
+			currentPrompt = fmt.Sprintf("%s\n\nPrevious attempt failed: %s\nPlease fix the output.", prompt, lastErr)
+		}
+
+		// Spawn a full sub-agent (RLM instance) for this item
+		subConfig := Config{
+			RecursiveModel: model,
+			APIBase:        am.apiBase,
+			APIKey:         am.apiKey,
+			MaxDepth:       config.MaxDepth,
+			MaxIterations:  config.MaxIter,
+			TimeoutSeconds: am.timeout,
+			ExtraParams:    am.extraParams,
+		}
+
+		subRLM := New(model, subConfig)
+		subRLM.currentDepth = 1 // Sub-agents start at depth 1
+		subRLM.observer = am.observer
+
+		// Build context with schema instructions if provided
+		context := ""
+		if config.OutputSchema != nil {
+			schemaJSON, _ := json.MarshalIndent(config.OutputSchema, "", "  ")
+			context = fmt.Sprintf("You must output valid JSON matching this schema:\n%s\n\nRespond with ONLY the JSON output.", string(schemaJSON))
+		}
+
+		result, stats, err := subRLM.Completion(currentPrompt, context)
+		subRLM.Shutdown()
+
+		tokensUsed := stats.TotalTokens
+
+		if err != nil {
+			lastErr = err.Error()
+			continue
+		}
+
+		// Extract JSON from the sub-agent's output
+		output := extractJSON(result)
+		if output == nil {
+			// Try wrapping the raw result as a string value
+			wrapped, _ := json.Marshal(result)
+			output = wrapped
+		}
+
+		// Validate against schema if provided
+		if config.OutputSchema != nil && output != nil {
+			if validationErr := validateMapOutput(output, config.OutputSchema); validationErr != "" {
+				lastErr = validationErr
+				continue
+			}
+		}
+
+		return agenticItemResult{
+			status:     MapItemCompleted,
+			output:     output,
+			retries:    attempt,
+			tokensUsed: tokensUsed,
+			llmCalls:   stats.LlmCalls,
+			iterations: stats.Iterations,
+		}
+	}
+
+	return agenticItemResult{
+		status:  MapItemFailed,
+		errMsg:  lastErr,
+		retries: config.MaxRetries,
+	}
+}
+
+// ─── Output Writing ─────────────────────────────────────────────────────────
+
+func writeAgenticOutput(path string, results []AgenticItemResult) error {
+	f, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer func() { _ = f.Close() }()
+
+	w := bufio.NewWriter(f)
+	for _, r := range results {
+		if r.Status == MapItemCompleted && r.Output != nil {
+			if _, writeErr := w.Write(r.Output); writeErr != nil {
+				return writeErr
+			}
+			if _, writeErr := w.WriteString("\n"); writeErr != nil {
+				return writeErr
+			}
+		} else {
+			errRecord := map[string]interface{}{
+				"_error":      r.Error,
+				"_index":      r.Index,
+				"_llm_calls":  r.LLMCalls,
+				"_iterations": r.Iterations,
+			}
+			data, _ := json.Marshal(errRecord)
+			if _, writeErr := w.Write(data); writeErr != nil {
+				return writeErr
+			}
+			if _, writeErr := w.WriteString("\n"); writeErr != nil {
+				return writeErr
+			}
+		}
+	}
+
+	return w.Flush()
+}

package/go/rlm/lcm_context_loop.go

@@ -0,0 +1,309 @@
+package rlm
+
+import (
+	"fmt"
+	"sync"
+)
+
+// ─── LCM Context Control Loop ───────────────────────────────────────────────
+// Implements the dual-threshold context management from the LCM paper:
+// - Below τ_soft: no overhead (zero-cost continuity)
+// - τ_soft ≤ tokens < τ_hard: async compaction between turns
+// - tokens ≥ τ_hard: blocking compaction before next LLM call
+
+// LCMConfig configures the Lossless Context Management engine.
+type LCMConfig struct {
+	// Enabled activates LCM context management (default: false for backward compat)
+	Enabled bool `json:"enabled"`
+
+	// SoftThreshold is τ_soft: token count above which async compaction begins.
+	// Default: 70% of model limit.
+	SoftThreshold int `json:"soft_threshold,omitempty"`
+
+	// HardThreshold is τ_hard: token count above which blocking compaction occurs.
+	// Default: 90% of model limit.
+	HardThreshold int `json:"hard_threshold,omitempty"`
+
+	// CompactionBlockSize is how many messages to compact at once.
+	// Default: 10 messages.
+	CompactionBlockSize int `json:"compaction_block_size,omitempty"`
+
+	// SummaryTargetTokens is the target size for each summary node.
+	// Default: 500 tokens.
+	SummaryTargetTokens int `json:"summary_target_tokens,omitempty"`
+}
+
+// DefaultLCMConfig returns default LCM configuration.
+func DefaultLCMConfig() LCMConfig {
+	return LCMConfig{
+		Enabled:             false,
+		CompactionBlockSize: 10,
+		SummaryTargetTokens: 500,
+	}
+}
+
+// LCMEngine is the main LCM context management engine.
+// It wraps the store, summarizer, and context control loop.
+type LCMEngine struct {
+	config     LCMConfig
+	store      *LCMStore
+	summarizer *LCMSummarizer
+	observer   *Observer
+	modelLimit int
+
+	// Async compaction state
+	compactMu     sync.Mutex
+	compacting    bool
+	compactResult chan *compactionResult
+}
+
+type compactionResult struct {
+	summary *SummaryNode
+	err     error
+}
+
+// NewLCMEngine creates a new LCM engine with the given configuration.
+func NewLCMEngine(config LCMConfig, store *LCMStore, summarizer *LCMSummarizer, observer *Observer, modelLimit int) *LCMEngine {
+	// Apply defaults based on model limit
+	if config.SoftThreshold == 0 && modelLimit > 0 {
+		config.SoftThreshold = int(float64(modelLimit) * 0.70)
+	}
+	if config.HardThreshold == 0 && modelLimit > 0 {
+		config.HardThreshold = int(float64(modelLimit) * 0.90)
+	}
+	if config.CompactionBlockSize == 0 {
+		config.CompactionBlockSize = 10
+	}
+	if config.SummaryTargetTokens == 0 {
+		config.SummaryTargetTokens = 500
+	}
+
+	return &LCMEngine{
+		config:     config,
+		store:      store,
+		summarizer: summarizer,
+		observer:   observer,
+		modelLimit: modelLimit,
+	}
+}
+
+// ─── Context Control Loop (Algorithm 2 from paper) ──────────────────────────
+
+// OnNewItem is called after each new message is added to the store.
+// It implements the context control loop from Figure 2 of the LCM paper.
+// Returns nil if no compaction was needed or if async compaction was triggered.
+func (e *LCMEngine) OnNewItem() error {
+	if !e.config.Enabled {
+		return nil
+	}
+
+	// Check if async compaction has completed
+	e.applyPendingCompaction()
+
+	tokens := e.store.ActiveContextTokens()
+
+	// Below soft threshold: zero-cost continuity
+	if tokens <= e.config.SoftThreshold {
+		return nil
+	}
+
+	// Soft threshold exceeded: trigger async compaction (non-blocking)
+	if tokens < e.config.HardThreshold {
+		e.observer.Debug("lcm.control", "Soft threshold exceeded (%d > %d), triggering async compaction",
+			tokens, e.config.SoftThreshold)
+		e.triggerAsyncCompaction()
+		return nil
+	}
+
+	// Hard threshold exceeded: blocking compaction
+	e.observer.Debug("lcm.control", "Hard threshold exceeded (%d >= %d), blocking compaction",
+		tokens, e.config.HardThreshold)
+	return e.blockingCompaction()
+}
+
+// ─── Async Compaction ───────────────────────────────────────────────────────
+
+func (e *LCMEngine) triggerAsyncCompaction() {
+	e.compactMu.Lock()
+	if e.compacting {
+		e.compactMu.Unlock()
+		return // Already compacting
+	}
+	e.compacting = true
+	e.compactResult = make(chan *compactionResult, 1)
+	e.compactMu.Unlock()
+
+	go func() {
+		result := e.performCompaction()
+		e.compactResult <- result
+	}()
+}
+
+func (e *LCMEngine) applyPendingCompaction() {
+	e.compactMu.Lock()
+	if !e.compacting || e.compactResult == nil {
+		e.compactMu.Unlock()
+		return
+	}
+
+	// Non-blocking check
+	select {
+	case result := <-e.compactResult:
+		e.compacting = false
+		e.compactMu.Unlock()
+
+		if result.err != nil {
+			e.observer.Error("lcm.control", "Async compaction failed: %v", result.err)
+			return
+		}
+		if result.summary != nil {
+			removed := e.store.CompactOldestBlock(result.summary)
+			e.observer.Debug("lcm.control", "Async compaction applied: replaced %d messages with summary %s",
+				len(removed), result.summary.ID)
+			e.observer.Event("lcm.compaction", map[string]string{
+				"type":           "async",
+				"summary_id":     result.summary.ID,
+				"messages_compacted": fmt.Sprintf("%d", len(removed)),
+				"summary_tokens": fmt.Sprintf("%d", result.summary.Tokens),
+				"level":          fmt.Sprintf("%d", result.summary.Level),
+			})
+		}
+	default:
+		e.compactMu.Unlock()
+		// Not done yet, continue
+	}
+}
+
+// ─── Blocking Compaction ────────────────────────────────────────────────────
+
+func (e *LCMEngine) blockingCompaction() error {
+	// Keep compacting until under hard threshold
+	for e.store.ActiveContextTokens() >= e.config.HardThreshold {
+		result := e.performCompaction()
+		if result.err != nil {
+			return fmt.Errorf("blocking compaction failed: %w", result.err)
+		}
+		if result.summary == nil {
+			break // Nothing more to compact
+		}
+
+		removed := e.store.CompactOldestBlock(result.summary)
+		e.observer.Debug("lcm.control", "Blocking compaction: replaced %d messages with summary %s (%d tokens)",
+			len(removed), result.summary.ID, result.summary.Tokens)
+		e.observer.Event("lcm.compaction", map[string]string{
+			"type":               "blocking",
+			"summary_id":         result.summary.ID,
+			"messages_compacted": fmt.Sprintf("%d", len(removed)),
+			"summary_tokens":     fmt.Sprintf("%d", result.summary.Tokens),
+			"level":              fmt.Sprintf("%d", result.summary.Level),
+		})
+	}
+	return nil
+}
+
+// ─── Core Compaction ────────────────────────────────────────────────────────
+
+func (e *LCMEngine) performCompaction() *compactionResult {
+	active := e.store.GetActiveContext()
+
+	// Find the oldest block of raw messages to compact (skip system prompt)
+	var block []*StoreMessage
+	for _, item := range active {
+		if item.IsMessage() {
+			if item.Message.Role == RoleSystem {
+				continue // Never compact system prompt
+			}
+			block = append(block, item.Message)
+			if len(block) >= e.config.CompactionBlockSize {
+				break
+			}
+		}
+	}
+
+	if len(block) == 0 {
+		return &compactionResult{summary: nil, err: nil}
+	}
+
+	// Apply three-level escalation
+	result, err := e.summarizer.SummarizeMessages(block, e.config.SummaryTargetTokens)
+	if err != nil {
+		return &compactionResult{err: err}
+	}
+
+	// Create summary node in the DAG
+	var msgIDs []string
+	for _, msg := range block {
+		msgIDs = append(msgIDs, msg.ID)
+	}
+
+	summary := e.store.CreateLeafSummary(msgIDs, result.Content, result.Level)
+
+	return &compactionResult{summary: summary}
+}
+
+// ─── Condensed Summaries (DAG depth > 1) ────────────────────────────────────
+
+// CondenseOldSummaries finds summary nodes in the active context and merges them
+// into a higher-order condensed summary. This creates DAG depth > 1.
+func (e *LCMEngine) CondenseOldSummaries() error {
+	active := e.store.GetActiveContext()
+
+	// Collect summary items
+	var summaryItems []*ActiveContextItem
+	for _, item := range active {
+		if !item.IsMessage() && item.Summary != nil {
+			summaryItems = append(summaryItems, item)
+		}
+	}
+
+	// Need at least 2 summaries to condense
+	if len(summaryItems) < 2 {
+		return nil
+	}
+
+	// Condense the oldest summaries
+	condenseCount := len(summaryItems)
+	if condenseCount > e.config.CompactionBlockSize {
+		condenseCount = e.config.CompactionBlockSize
+	}
+	toCondense := summaryItems[:condenseCount]
+
+	// Build combined content for re-summarization
+	var combined string
+	var childIDs []string
+	for _, item := range toCondense {
+		combined += item.Summary.Content + "\n\n"
+		childIDs = append(childIDs, item.Summary.ID)
+	}
+
+	// Summarize the combined summaries
+	result, err := e.summarizer.Summarize(combined, e.config.SummaryTargetTokens)
+	if err != nil {
+		return fmt.Errorf("condensation failed: %w", err)
+	}
+
+	// Create condensed summary node
+	condensed := e.store.CreateCondensedSummary(childIDs, result.Content, result.Level)
+
+	e.observer.Debug("lcm.control", "Condensed %d summaries into %s (%d tokens)",
+		len(childIDs), condensed.ID, condensed.Tokens)
+
+	return nil
+}
+
+// ─── Query Helpers ──────────────────────────────────────────────────────────
+
+// GetStore returns the underlying LCM store.
+func (e *LCMEngine) GetStore() *LCMStore {
+	return e.store
+}
+
+// GetConfig returns the LCM configuration.
+func (e *LCMEngine) GetConfig() LCMConfig {
+	return e.config
+}
+
+// IsEnabled returns whether LCM is active.
+func (e *LCMEngine) IsEnabled() bool {
+	return e.config.Enabled
+}