@phuetz/code-buddy 0.3.0 → 1.0.0-rc.8

package/dist/agent/reasoning/mcts.js

@@ -2,9 +2,25 @@
  * Monte Carlo Tree Search for Code Generation
  *
  * Implements MCTS with the Rethink mechanism for refining erroneous thoughts.
- * Based on RethinkMCTS (arXiv 2409.09584).
+ * Based on:
+ * - RethinkMCTS (arXiv 2409.09584)
+ * - MCTSr Q-value formula (arXiv 2406.07394)
+ * - BFS beam search for Tree-of-Thought
  */
-import { DEFAULT_MCTS_CONFIG, } from "./types.js";
+import { DEFAULT_MCTS_CONFIG, THINKING_MODE_CONFIG, } from './types.js';
+/**
+ * Estimated tokens per LLM call type
+ */
+const TOKENS_PER_GENERATE = 500;
+const TOKENS_PER_EVALUATE = 200;
+/**
+ * Categorical evaluation labels mapped to numeric scores
+ */
+const CATEGORICAL_SCORES = {
+    'sure': 1.0,
+    'likely': 0.6,
+    'impossible': 0.1,
+};
 /**
  * Create a unique node ID
  */
@@ -38,58 +54,207 @@ export class MCTS {
             maxDepthReached: 0,
             totalTime: 0,
             bestScore: 0,
+            tokensUsed: 0,
         };
     }
     /**
-     * Run MCTS search for a problem
+     * Public search router — dispatches to MCTS or BFS based on config
      */
     async search(problem) {
+        if (this.config.progressiveDeepening) {
+            return this.searchProgressive(problem);
+        }
+        if (this.config.searchAlgorithm === 'bfs') {
+            return this.searchBFS(problem);
+        }
+        return this.searchMCTS(problem);
+    }
+    /**
+     * Run classic MCTS search for a problem
+     */
+    async searchMCTS(problem) {
         this.startTime = Date.now();
-        this.stats = {
-            iterations: 0,
-            nodesCreated: 0,
-            nodesEvaluated: 0,
-            nodesRefined: 0,
-            maxDepthReached: 0,
-            totalTime: 0,
-            bestScore: 0,
-        };
+        this.resetStats();
         // Create root node
-        this.root = this.createNode(`Understanding the problem: ${problem.description}`, "analysis", null, 0);
+        this.root = this.createNode(`Understanding the problem: ${problem.description}`, 'analysis', null, 0);
         // Run MCTS iterations
         for (let i = 0; i < this.config.maxIterations; i++) {
             // Check time limit
             if (this.config.timeLimit && Date.now() - this.startTime > this.config.timeLimit) {
                 break;
             }
+            // Check token budget
+            if (this.isTokenBudgetExhausted()) {
+                break;
+            }
             this.stats.iterations = i + 1;
             // 1. Selection
             const selectedNode = this.select(this.root);
             // 2. Expansion
             if (selectedNode.depth < this.config.maxDepth) {
                 await this.expand(selectedNode, problem);
+                this.emitProgress('expansion', selectedNode.id);
             }
             // 3. Simulation & Evaluation
             if (selectedNode.children.length > 0) {
                 for (const child of selectedNode.children) {
+                    if (this.isTokenBudgetExhausted())
+                        break;
                     await this.simulate(child, problem);
+                    this.emitProgress('evaluation', child.id);
                 }
             }
-            // 4. Backpropagation
+            // 4. Backpropagation (MCTSr formula)
             this.backpropagate(selectedNode);
             // 5. Rethink (if enabled)
             if (this.config.useRethink) {
                 await this.rethink(selectedNode, problem);
             }
+            this.emitProgress('iteration');
             // Check for solution
             const solution = this.findBestSolution();
             if (solution && solution.score > 0.9) {
+                this.emitProgress('solution_found', solution.id);
+                break;
+            }
+        }
+        this.stats.totalTime = Date.now() - this.startTime;
+        return this.buildResult();
+    }
+    /**
+     * BFS Beam Search mode
+     *
+     * At each depth level, keep top-b states (beamWidth) and discard the rest.
+     * Uses categorical evaluation ("sure/likely/impossible") instead of numeric.
+     */
+    async searchBFS(problem) {
+        this.startTime = Date.now();
+        this.resetStats();
+        const beamWidth = this.config.beamWidth ?? 3;
+        // Create root node
+        this.root = this.createNode(`Understanding the problem: ${problem.description}`, 'analysis', null, 0);
+        let currentBeam = [this.root];
+        for (let depth = 0; depth < this.config.maxDepth; depth++) {
+            if (this.config.timeLimit && Date.now() - this.startTime > this.config.timeLimit) {
+                break;
+            }
+            if (this.isTokenBudgetExhausted()) {
+                break;
+            }
+            const candidates = [];
+            // Expand all nodes in the current beam
+            for (const node of currentBeam) {
+                if (this.isTokenBudgetExhausted())
+                    break;
+                const thoughts = await this.generateThoughts(node, problem);
+                this.stats.tokensUsed += TOKENS_PER_GENERATE;
+                for (const thought of thoughts.slice(0, this.config.expansionCount)) {
+                    const childType = this.determineThoughtType(thought, depth + 1);
+                    const child = this.createNode(thought, childType, node, depth + 1);
+                    node.children.push(child);
+                    candidates.push(child);
+                    this.stats.maxDepthReached = Math.max(this.stats.maxDepthReached, child.depth);
+                }
+                node.state = 'evaluated';
+                this.emitProgress('expansion', node.id);
+            }
+            if (candidates.length === 0)
+                break;
+            // Evaluate all candidates using categorical evaluation
+            for (const candidate of candidates) {
+                if (this.isTokenBudgetExhausted())
+                    break;
+                const rawScore = await this.evaluateThought(candidate, problem);
+                this.stats.tokensUsed += TOKENS_PER_EVALUATE;
+                // Map to categorical: nearest categorical label
+                candidate.score = this.toCategoricalScore(rawScore);
+                candidate.visits = 1;
+                candidate.metadata.rewardSamples = [candidate.score];
+                this.stats.nodesEvaluated++;
+                // If implementation, try execution
+                if (candidate.type === 'implementation') {
+                    const code = this.extractCode(candidate.content);
+                    if (code) {
+                        const result = await this.executeCode(code);
+                        candidate.metadata.codeGenerated = code;
+                        candidate.metadata.executionResult = result;
+                        if (result.success) {
+                            candidate.score = Math.min(1, candidate.score + 0.3);
+                            candidate.state = 'completed';
+                        }
+                        else {
+                            candidate.score = Math.max(0, candidate.score - 0.2);
+                            candidate.metadata.feedback = result.error;
+                        }
+                    }
+                }
+                this.stats.bestScore = Math.max(this.stats.bestScore, candidate.score);
+                this.emitProgress('evaluation', candidate.id);
+            }
+            // Keep top-b candidates by score
+            candidates.sort((a, b) => b.score - a.score);
+            const kept = candidates.slice(0, beamWidth);
+            const pruned = candidates.slice(beamWidth);
+            for (const p of pruned) {
+                p.state = 'pruned';
+            }
+            currentBeam = kept;
+            this.stats.iterations++;
+            this.emitProgress('iteration');
+            // Check for solution
+            const best = kept[0];
+            if (best && best.score > 0.9) {
+                this.emitProgress('solution_found', best.id);
                 break;
             }
         }
         this.stats.totalTime = Date.now() - this.startTime;
         return this.buildResult();
     }
+    /**
+     * Progressive deepening search
+     *
+     * Starts with shallow config (5 iterations). If bestScore < 0.6,
+     * escalates to medium, then deep, then exhaustive.
+     */
+    async searchProgressive(problem) {
+        const escalationOrder = [
+            'shallow', 'medium', 'deep', 'exhaustive',
+        ];
+        const scoreThreshold = 0.6;
+        let lastResult = null;
+        for (const mode of escalationOrder) {
+            const modeConfig = THINKING_MODE_CONFIG[mode];
+            // Apply mode config but keep callbacks and overrides from original config
+            this.config = {
+                ...DEFAULT_MCTS_CONFIG,
+                ...modeConfig,
+                // Preserve user overrides for non-mode fields
+                searchAlgorithm: this.config.searchAlgorithm,
+                useRethink: this.config.useRethink,
+                rethinkThreshold: this.config.rethinkThreshold,
+                timeLimit: this.config.timeLimit,
+                rewardSamples: this.config.rewardSamples,
+                onProgress: this.config.onProgress,
+                // Use mode token budget if no explicit override
+                tokenBudget: this.config.tokenBudget ?? modeConfig.tokenBudget,
+                // Disable progressive deepening to prevent recursion
+                progressiveDeepening: false,
+            };
+            if (this.config.searchAlgorithm === 'bfs') {
+                lastResult = await this.searchBFS(problem);
+            }
+            else {
+                lastResult = await this.searchMCTS(problem);
+            }
+            if (lastResult.stats.bestScore >= scoreThreshold) {
+                break;
+            }
+            // If we haven't reached the threshold, accumulate stats for next round
+            // The next round starts fresh but benefits from higher iteration counts
+        }
+        return lastResult;
+    }
     /**
      * Create a thought node
      */
@@ -106,8 +271,9 @@ export class MCTS {
             depth,
             metadata: {
                 generationRound: this.stats.iterations,
+                rewardSamples: [],
             },
-            state: "pending",
+            state: 'pending',
         };
     }
     /**
@@ -122,7 +288,7 @@ export class MCTS {
         let bestChild = null;
         let bestUCB = -Infinity;
         for (const child of node.children) {
-            if (child.state === "pruned")
+            if (child.state === 'pruned')
                 continue;
             const ucb = this.calculateUCB1(child, node.visits);
             if (ucb > bestUCB) {
@@ -138,23 +304,40 @@ export class MCTS {
     }
     /**
      * Calculate UCB1 value for a node
+     * Uses MCTSr Q-value: Q(a) = 0.5 * (min(R_a) + mean(R_a))
      */
     calculateUCB1(node, parentVisits) {
         if (node.visits === 0) {
             return Infinity; // Encourage exploration of unvisited nodes
         }
-        const exploitation = node.score / node.visits;
+        const qValue = this.computeQValue(node);
         const exploration = this.config.explorationConstant *
             Math.sqrt(Math.log(parentVisits) / node.visits);
-        return exploitation + exploration;
+        return qValue + exploration;
+    }
+    /**
+     * Compute MCTSr Q-value for a node
+     * Q(a) = 0.5 * (min(R_a) + mean(R_a))
+     * Falls back to simple score/visits if no reward samples
+     */
+    computeQValue(node) {
+        const samples = node.metadata.rewardSamples;
+        if (!samples || samples.length === 0) {
+            // node.score is already a Q-value in [0,1], not cumulative — don't divide
+            return node.score;
+        }
+        const minReward = Math.min(...samples);
+        const meanReward = samples.reduce((sum, r) => sum + r, 0) / samples.length;
+        return 0.5 * (minReward + meanReward);
     }
     /**
      * Expand a node by generating child thoughts
      */
     async expand(node, problem) {
-        node.state = "exploring";
+        node.state = 'exploring';
         // Generate new thoughts
         const thoughts = await this.generateThoughts(node, problem);
+        this.stats.tokensUsed += TOKENS_PER_GENERATE;
         // Create child nodes
         for (const thought of thoughts.slice(0, this.config.expansionCount)) {
             const childType = this.determineThoughtType(thought, node.depth);
@@ -162,45 +345,57 @@ export class MCTS {
             node.children.push(child);
             this.stats.maxDepthReached = Math.max(this.stats.maxDepthReached, child.depth);
         }
-        node.state = "evaluated";
+        node.state = 'evaluated';
     }
     /**
      * Determine the type of thought based on content and depth
      */
     determineThoughtType(content, depth) {
         const lowerContent = content.toLowerCase();
-        if (lowerContent.includes("```") || lowerContent.includes("function") ||
-            lowerContent.includes("class") || lowerContent.includes("def ")) {
-            return "implementation";
+        if (lowerContent.includes('```') || lowerContent.includes('function') ||
+            lowerContent.includes('class') || lowerContent.includes('def ')) {
+            return 'implementation';
         }
-        if (lowerContent.includes("test") || lowerContent.includes("verify") ||
-            lowerContent.includes("check")) {
-            return "verification";
+        if (lowerContent.includes('test') || lowerContent.includes('verify') ||
+            lowerContent.includes('check')) {
+            return 'verification';
         }
-        if (lowerContent.includes("improve") || lowerContent.includes("refine") ||
-            lowerContent.includes("optimize")) {
-            return "refinement";
+        if (lowerContent.includes('improve') || lowerContent.includes('refine') ||
+            lowerContent.includes('optimize')) {
+            return 'refinement';
         }
-        if (lowerContent.includes("therefore") || lowerContent.includes("conclusion") ||
-            lowerContent.includes("solution")) {
-            return "conclusion";
+        if (lowerContent.includes('therefore') || lowerContent.includes('conclusion') ||
+            lowerContent.includes('solution')) {
+            return 'conclusion';
         }
         if (depth <= 1) {
-            return "analysis";
+            return 'analysis';
         }
-        return "hypothesis";
+        return 'hypothesis';
     }
     /**
      * Simulate from a node (rollout)
+     * Collects multiple reward samples per node for robust Q-value estimation
      */
     async simulate(node, problem) {
-        // Evaluate the thought
-        const score = await this.evaluateThought(node, problem);
-        node.score = score;
+        const numSamples = this.config.rewardSamples ?? 1;
+        if (!node.metadata.rewardSamples) {
+            node.metadata.rewardSamples = [];
+        }
+        // Collect reward samples
+        for (let s = 0; s < numSamples; s++) {
+            if (this.isTokenBudgetExhausted())
+                break;
+            const score = await this.evaluateThought(node, problem);
+            this.stats.tokensUsed += TOKENS_PER_EVALUATE;
+            node.metadata.rewardSamples.push(score);
+        }
+        // Set node score as MCTSr Q-value
+        node.score = this.computeQValue(node);
         node.visits = 1;
         this.stats.nodesEvaluated++;
         // If implementation, try to execute
-        if (node.type === "implementation") {
+        if (node.type === 'implementation') {
             const code = this.extractCode(node.content);
             if (code) {
                 const result = await this.executeCode(code);
@@ -209,7 +404,7 @@ export class MCTS {
                 // Adjust score based on execution
                 if (result.success) {
                     node.score = Math.min(1, node.score + 0.3);
-                    node.state = "completed";
+                    node.state = 'completed';
                 }
                 else {
                     node.score = Math.max(0, node.score - 0.2);
@@ -220,19 +415,23 @@ export class MCTS {
         this.stats.bestScore = Math.max(this.stats.bestScore, node.score);
     }
     /**
-     * Backpropagate scores up the tree
+     * Backpropagate scores up the tree using MCTSr formula
+     * Q'(parent) = 0.5 * (Q(parent) + max(Q(children)))
      */
     backpropagate(node) {
         let current = node;
         while (current !== null) {
             current.visits++;
-            // Update score as average of children
+            // Update score using MCTSr backpropagation formula
             if (current.children.length > 0) {
-                const childScores = current.children
-                    .filter(c => c.state !== "pruned")
-                    .map(c => c.score);
-                if (childScores.length > 0) {
-                    current.score = Math.max(...childScores);
+                const childQValues = current.children
+                    .filter(c => c.state !== 'pruned')
+                    .map(c => this.computeQValue(c));
+                if (childQValues.length > 0) {
+                    const maxChildQ = Math.max(...childQValues);
+                    const parentQ = this.computeQValue(current);
+                    // MCTSr: Q'(parent) = 0.5 * (Q(parent) + max(Q(children)))
+                    current.score = 0.5 * (parentQ + maxChildQ);
                 }
             }
             current = current.parent;
@@ -245,18 +444,22 @@ export class MCTS {
         // Find nodes that failed execution or have low scores
         const nodesToRethink = this.findNodesNeedingRethink(node);
         for (const n of nodesToRethink) {
+            if (this.isTokenBudgetExhausted())
+                break;
             if (n.metadata.feedback) {
                 const refinedContent = await this.refineThought(n, n.metadata.feedback);
+                this.stats.tokensUsed += TOKENS_PER_GENERATE;
                 // Create refined node as sibling
                 const refinedNode = this.createNode(refinedContent, n.type, n.parent, n.depth);
                 refinedNode.metadata.reasoning = `Refined from: ${n.id}`;
-                refinedNode.state = "refined";
+                refinedNode.state = 'refined';
                 if (n.parent) {
                     n.parent.children.push(refinedNode);
                 }
                 // Mark original as pruned
-                n.state = "pruned";
+                n.state = 'pruned';
                 this.stats.nodesRefined++;
+                this.emitProgress('refinement', refinedNode.id);
             }
         }
     }
@@ -266,7 +469,7 @@ export class MCTS {
     findNodesNeedingRethink(node) {
         const result = [];
         const traverse = (n) => {
-            if (n.state !== "pruned" &&
+            if (n.state !== 'pruned' &&
                 n.score < this.config.rethinkThreshold &&
                 n.metadata.feedback) {
                 result.push(n);
@@ -287,8 +490,8 @@ export class MCTS {
         let best = null;
         let bestScore = -1;
         const traverse = (node) => {
-            if (node.state !== "pruned" &&
-                (node.type === "implementation" || node.type === "conclusion") &&
+            if (node.state !== 'pruned' &&
+                (node.type === 'implementation' || node.type === 'conclusion') &&
                 node.score > bestScore) {
                 bestScore = node.score;
                 best = node;
@@ -321,8 +524,8 @@ export class MCTS {
         const alternatives = [];
         const traverse = (node) => {
             if (node !== best &&
-                node.state !== "pruned" &&
-                (node.type === "implementation" || node.type === "conclusion") &&
+                node.state !== 'pruned' &&
+                (node.type === 'implementation' || node.type === 'conclusion') &&
                 node.score > 0.5) {
                 alternatives.push(node);
             }
@@ -341,12 +544,13 @@ export class MCTS {
     extractCode(content) {
         // Try to extract code blocks
         const codeBlockMatch = content.match(/```(?:\w+)?\n([\s\S]*?)```/);
-        if (codeBlockMatch) {
-            return codeBlockMatch[1].trim();
+        const captured = codeBlockMatch?.[1];
+        if (captured !== undefined) {
+            return captured.trim();
         }
         // Check if content looks like code
-        if (content.includes("function") || content.includes("class") ||
-            content.includes("const") || content.includes("def ")) {
+        if (content.includes('function') || content.includes('class') ||
+            content.includes('const') || content.includes('def ')) {
             return content;
         }
         return null;
@@ -367,6 +571,66 @@ export class MCTS {
             tree: this.root,
         };
     }
+    /**
+     * Reset stats for a new search
+     */
+    resetStats() {
+        this.stats = {
+            iterations: 0,
+            nodesCreated: 0,
+            nodesEvaluated: 0,
+            nodesRefined: 0,
+            maxDepthReached: 0,
+            totalTime: 0,
+            bestScore: 0,
+            tokensUsed: 0,
+        };
+    }
+    /**
+     * Check if the token budget has been exhausted
+     */
+    isTokenBudgetExhausted() {
+        if (!this.config.tokenBudget)
+            return false;
+        return this.stats.tokensUsed >= this.config.tokenBudget;
+    }
+    /**
+     * Map a numeric score to the nearest categorical score
+     * Categories: "sure" (1.0), "likely" (0.6), "impossible" (0.1)
+     */
+    toCategoricalScore(raw) {
+        const entries = Object.entries(CATEGORICAL_SCORES);
+        const first = entries[0];
+        if (first === undefined)
+            return raw;
+        let closest = first[1];
+        let minDist = Math.abs(raw - closest);
+        for (const [, value] of entries) {
+            const dist = Math.abs(raw - value);
+            if (dist < minDist) {
+                minDist = dist;
+                closest = value;
+            }
+        }
+        return closest;
+    }
+    /**
+     * Emit a progress event if a callback is configured
+     */
+    emitProgress(type, nodeId, message) {
+        if (!this.config.onProgress)
+            return;
+        this.config.onProgress({
+            type,
+            iteration: this.stats.iterations,
+            nodesCreated: this.stats.nodesCreated,
+            nodesEvaluated: this.stats.nodesEvaluated,
+            bestScore: this.stats.bestScore,
+            tokensUsed: this.stats.tokensUsed,
+            nodeId,
+            message,
+        });
+    }
     /**
      * Get current statistics
      */
@@ -382,31 +646,31 @@ export class MCTS {
     /**
      * Format tree for display
      */
-    formatTree(node = this.root, indent = "") {
+    formatTree(node = this.root, indent = '') {
         if (!node)
-            return "Empty tree";
+            return 'Empty tree';
         const stateEmoji = {
-            pending: "⏳",
-            exploring: "🔍",
-            evaluated: "📊",
-            refined: "🔄",
-            completed: "✅",
-            failed: "❌",
-            pruned: "✂️",
+            pending: '⏳',
+            exploring: '🔍',
+            evaluated: '📊',
+            refined: '🔄',
+            completed: '✅',
+            failed: '❌',
+            pruned: '✂️',
         };
         const typeEmoji = {
-            analysis: "🔬",
-            hypothesis: "💡",
-            implementation: "💻",
-            verification: "✔️",
-            refinement: "🔧",
-            conclusion: "🎯",
+            analysis: '🔬',
+            hypothesis: '💡',
+            implementation: '💻',
+            verification: '✔️',
+            refinement: '🔧',
+            conclusion: '🎯',
         };
         let output = `${indent}${stateEmoji[node.state]} ${typeEmoji[node.type]} `;
         output += `[${node.score.toFixed(2)}, v:${node.visits}] `;
-        output += `${node.content.slice(0, 50)}${node.content.length > 50 ? "..." : ""}\n`;
+        output += `${node.content.slice(0, 50)}${node.content.length > 50 ? '...' : ''}\n`;
         for (const child of node.children) {
-            output += this.formatTree(child, indent + "  ");
+            output += this.formatTree(child, indent + '  ');
         }
         return output;
     }