recursive-llm-ts 4.9.0 → 5.0.0

package/go/rlm/lcm_map.go

@@ -0,0 +1,348 @@
+package rlm
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// ─── LLM-Map Operator ──────────────────────────────────────────────────────
+// Implements Operator-Level Recursion from the LCM paper (Section 3.1).
+// Processes each item in a JSONL file via independent LLM calls with
+// schema validation, retry logic, and database-backed execution tracking.
+
+// LLMMapConfig configures an LLM-Map operation.
+type LLMMapConfig 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 validation
+	Concurrency  int         `json:"concurrency"`   // Worker pool size (default: 16)
+	MaxRetries   int         `json:"max_retries"`   // Per-item retry limit (default: 3)
+	Model        string      `json:"model"`         // Model to use (default: engine model)
+}
+
+// LLMMapResult contains the results of an LLM-Map operation.
+type LLMMapResult 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   []MapItemResult `json:"item_results,omitempty"`
+}
+
+// MapItemResult tracks the status of a single item in the map operation.
+type MapItemResult struct {
+	Index   int             `json:"index"`
+	Status  MapItemStatus   `json:"status"`
+	Output  json.RawMessage `json:"output,omitempty"`
+	Error   string          `json:"error,omitempty"`
+	Retries int             `json:"retries"`
+}
+
+// MapItemStatus represents the execution status of a map item.
+type MapItemStatus string
+
+const (
+	MapItemPending   MapItemStatus = "pending"
+	MapItemRunning   MapItemStatus = "running"
+	MapItemCompleted MapItemStatus = "completed"
+	MapItemFailed    MapItemStatus = "failed"
+)
+
+// LLMMapper executes LLM-Map operations.
+type LLMMapper struct {
+	model       string
+	apiBase     string
+	apiKey      string
+	timeout     int
+	extraParams map[string]interface{}
+	observer    *Observer
+}
+
+// NewLLMMapper creates a new LLM-Map executor.
+func NewLLMMapper(model, apiBase, apiKey string, timeout int, extraParams map[string]interface{}, observer *Observer) *LLMMapper {
+	return &LLMMapper{
+		model:       model,
+		apiBase:     apiBase,
+		apiKey:      apiKey,
+		timeout:     timeout,
+		extraParams: extraParams,
+		observer:    observer,
+	}
+}
+
+// Execute runs an LLM-Map operation: parallel LLM calls over JSONL input.
+func (m *LLMMapper) Execute(config LLMMapConfig) (*LLMMapResult, error) {
+	start := time.Now()
+
+	// Apply defaults
+	if config.Concurrency <= 0 {
+		config.Concurrency = 16
+	}
+	if config.MaxRetries <= 0 {
+		config.MaxRetries = 3
+	}
+	model := config.Model
+	if model == "" {
+		model = m.model
+	}
+
+	m.observer.Debug("lcm.map", "Starting LLM-Map: input=%s, concurrency=%d, model=%s",
+		config.InputPath, config.Concurrency, model)
+
+	// Read input items
+	items, err := readJSONLFile(config.InputPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read input: %w", err)
+	}
+
+	m.observer.Debug("lcm.map", "Read %d items from %s", len(items), config.InputPath)
+
+	// Initialize results tracking
+	results := make([]MapItemResult, len(items))
+	for i := range results {
+		results[i] = MapItemResult{
+			Index:  i,
+			Status: MapItemPending,
+		}
+	}
+
+	// Worker pool
+	var wg sync.WaitGroup
+	itemChan := make(chan int, len(items))
+	var totalTokens int64
+
+	// Feed items to workers
+	for i := range items {
+		itemChan <- i
+	}
+	close(itemChan)
+
+	// Spawn workers
+	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 := m.processItem(items[idx], config.Prompt, config.OutputSchema, model, config.MaxRetries)
+				atomic.AddInt64(&totalTokens, int64(result.tokensUsed))
+
+				mu.Lock()
+				results[idx] = MapItemResult{
+					Index:   idx,
+					Status:  result.status,
+					Output:  result.output,
+					Error:   result.errMsg,
+					Retries: result.retries,
+				}
+				mu.Unlock()
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	// Write output file
+	if config.OutputPath != "" {
+		if err := writeJSONLOutput(config.OutputPath, results); err != nil {
+			return nil, fmt.Errorf("failed to write output: %w", err)
+		}
+	}
+
+	// Count results
+	completed, failed := 0, 0
+	for _, r := range results {
+		switch r.Status {
+		case MapItemCompleted:
+			completed++
+		case MapItemFailed:
+			failed++
+		}
+	}
+
+	duration := time.Since(start)
+	m.observer.Debug("lcm.map", "LLM-Map complete: %d/%d succeeded, %d failed in %s",
+		completed, len(items), failed, duration)
+	m.observer.Event("lcm.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 &LLMMapResult{
+		TotalItems:  len(items),
+		Completed:   completed,
+		Failed:      failed,
+		OutputPath:  config.OutputPath,
+		Duration:    duration,
+		TokensUsed:  int(totalTokens),
+		ItemResults: results,
+	}, nil
+}
+
+// ─── Per-Item Processing ────────────────────────────────────────────────────
+
+type itemProcessResult struct {
+	status     MapItemStatus
+	output     json.RawMessage
+	errMsg     string
+	retries    int
+	tokensUsed int
+}
+
+func (m *LLMMapper) processItem(item json.RawMessage, promptTemplate string, schema *JSONSchema, model string, maxRetries int) itemProcessResult {
+	// Build prompt by replacing {{item}} placeholder
+	prompt := strings.ReplaceAll(promptTemplate, "{{item}}", string(item))
+
+	var lastErr string
+	for attempt := 0; attempt <= maxRetries; attempt++ {
+		// Add validation feedback on retry
+		if attempt > 0 && lastErr != "" {
+			prompt = fmt.Sprintf("%s\n\nPrevious attempt failed validation: %s\nPlease fix the output to match the required schema.", prompt, lastErr)
+		}
+
+		request := ChatRequest{
+			Model: model,
+			Messages: []Message{
+				{Role: "user", Content: prompt},
+			},
+			APIBase:     m.apiBase,
+			APIKey:      m.apiKey,
+			Timeout:     m.timeout,
+			ExtraParams: m.extraParams,
+		}
+
+		result, err := CallChatCompletion(request)
+		if err != nil {
+			lastErr = err.Error()
+			continue
+		}
+
+		tokensUsed := 0
+		if result.Usage != nil {
+			tokensUsed = result.Usage.TotalTokens
+		}
+
+		// Extract JSON from response
+		output := extractJSON(result.Content)
+		if output == nil {
+			lastErr = "no valid JSON found in response"
+			continue
+		}
+
+		// Validate against schema if provided
+		if schema != nil {
+			if validationErr := validateMapOutput(output, schema); validationErr != "" {
+				lastErr = validationErr
+				continue
+			}
+		}
+
+		return itemProcessResult{
+			status:     MapItemCompleted,
+			output:     output,
+			retries:    attempt,
+			tokensUsed: tokensUsed,
+		}
+	}
+
+	return itemProcessResult{
+		status:  MapItemFailed,
+		errMsg:  lastErr,
+		retries: maxRetries,
+	}
+}
+
+// ─── JSONL I/O ──────────────────────────────────────────────────────────────
+
+func readJSONLFile(path string) ([]json.RawMessage, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = f.Close() }()
+
+	var items []json.RawMessage
+	scanner := bufio.NewScanner(f)
+	// Increase buffer for large lines
+	scanner.Buffer(make([]byte, 0, 1024*1024), 10*1024*1024)
+
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" {
+			continue
+		}
+		items = append(items, json.RawMessage(line))
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("error reading JSONL: %w", err)
+	}
+
+	return items, nil
+}
+
+func writeJSONLOutput(path string, results []MapItemResult) 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 {
+			// Write error record
+			errRecord := map[string]interface{}{
+				"_error": r.Error,
+				"_index": r.Index,
+			}
+			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()
+}
+
+// ─── Helpers ────────────────────────────────────────────────────────────────
+
+// extractJSON delegates to the shared ExtractFirstJSON (json_extraction.go).
+var extractJSON = ExtractFirstJSON
+
+// validateMapOutput validates JSON output against a schema.
+func validateMapOutput(output json.RawMessage, schema *JSONSchema) string {
+	var parsed interface{}
+	if err := json.Unmarshal(output, &parsed); err != nil {
+		return fmt.Sprintf("invalid JSON: %v", err)
+	}
+
+	// Use the existing schema validation infrastructure
+	if err := validateValue(parsed, schema); err != nil {
+		return err.Error()
+	}
+	return ""
+}