Tree of Thoughts

1from dataclasses import dataclass, field
2from typing import Optional, Callable
3import heapq
4
5@dataclass
6class ThoughtNode:
7    """A node in the tree of thoughts."""
8    thought: str
9    state: dict  # Problem state after this thought
10    parent: Optional['ThoughtNode'] = None
11    children: list['ThoughtNode'] = field(default_factory=list)
12    score: float = 0.0
13    depth: int = 0
14
15    def get_path(self) -> list[str]:
16        """Get the thought sequence from root to this node."""
17        path = []
18        node = self
19        while node:
20            path.append(node.thought)
21            node = node.parent
22        return list(reversed(path))
23
24    def __lt__(self, other: 'ThoughtNode') -> bool:
25        """For priority queue comparison."""
26        return self.score > other.score  # Higher score = higher priority
27
28class TreeOfThoughts:
29    """
30    Tree of Thoughts reasoning framework.
31
32    Explores multiple reasoning paths, evaluates partial
33    solutions, and finds optimal thought sequences.
34    """
35
36    def __init__(
37        self,
38        thought_generator: Callable,
39        state_evaluator: Callable,
40        max_depth: int = 10,
41        branching_factor: int = 3
42    ):
43        self.generate_thoughts = thought_generator
44        self.evaluate_state = state_evaluator
45        self.max_depth = max_depth
46        self.branching_factor = branching_factor
47
48    def solve(
49        self,
50        problem: str,
51        initial_state: dict,
52        search_method: str = "best_first"
53    ) -> Optional[ThoughtNode]:
54        """Solve problem using tree search."""
55
56        root = ThoughtNode(
57            thought=f"Problem: {problem}",
58            state=initial_state,
59            depth=0
60        )
61        root.score = self.evaluate_state(root.state)
62
63        if search_method == "bfs":
64            return self._bfs_search(root)
65        elif search_method == "dfs":
66            return self._dfs_search(root)
67        elif search_method == "best_first":
68            return self._best_first_search(root)
69        else:
70            raise ValueError(f"Unknown search method: {search_method}")
71
72    def _expand_node(self, node: ThoughtNode) -> list[ThoughtNode]:
73        """Generate child nodes for a node."""
74        if node.depth >= self.max_depth:
75            return []
76
77        # Generate candidate thoughts
78        thoughts = self.generate_thoughts(
79            node.state,
80            node.get_path(),
81            self.branching_factor
82        )
83
84        children = []
85        for thought, new_state in thoughts:
86            child = ThoughtNode(
87                thought=thought,
88                state=new_state,
89                parent=node,
90                depth=node.depth + 1
91            )
92            child.score = self.evaluate_state(new_state)
93            children.append(child)
94            node.children.append(child)
95
96        return children

1def _bfs_search(self, root: ThoughtNode) -> Optional[ThoughtNode]:
2    """Breadth-first tree search."""
3    from collections import deque
4
5    queue = deque([root])
6    best_solution = None
7    best_score = float('-inf')
8
9    while queue:
10        node = queue.popleft()
11
12        # Check if this is a complete solution
13        if self._is_solution(node.state):
14            if node.score > best_score:
15                best_solution = node
16                best_score = node.score
17            continue
18
19        # Skip if score too low (pruning)
20        if node.score < best_score * 0.5:
21            continue
22
23        # Expand children
24        children = self._expand_node(node)
25
26        # Sort children by score and keep top k
27        children.sort(key=lambda n: n.score, reverse=True)
28        for child in children[:self.branching_factor]:
29            queue.append(child)
30
31    return best_solution

1def _dfs_search(
2    self,
3    root: ThoughtNode,
4    prune_threshold: float = 0.3
5) -> Optional[ThoughtNode]:
6    """Depth-first tree search with pruning."""
7
8    stack = [root]
9    best_solution = None
10    best_score = float('-inf')
11
12    while stack:
13        node = stack.pop()
14
15        # Check solution
16        if self._is_solution(node.state):
17            if node.score > best_score:
18                best_solution = node
19                best_score = node.score
20            continue
21
22        # Prune unpromising branches
23        if node.score < prune_threshold:
24            continue
25
26        # Expand and add children to stack
27        children = self._expand_node(node)
28
29        # Add in reverse order so best is explored first
30        children.sort(key=lambda n: n.score)
31        for child in children:
32            stack.append(child)
33
34    return best_solution

1def _best_first_search(self, root: ThoughtNode) -> Optional[ThoughtNode]:
2    """Best-first tree search using priority queue."""
3
4    # Priority queue: (negative_score, node) - heapq is min-heap
5    frontier = []
6    heapq.heappush(frontier, (-root.score, id(root), root))
7
8    best_solution = None
9    best_score = float('-inf')
10    explored = 0
11    max_explored = 1000  # Limit exploration
12
13    while frontier and explored < max_explored:
14        _, _, node = heapq.heappop(frontier)
15        explored += 1
16
17        # Check solution
18        if self._is_solution(node.state):
19            if node.score > best_score:
20                best_solution = node
21                best_score = node.score
22                # Could return immediately for first solution
23                # or continue to find better
24            continue
25
26        # Skip if we've found a great solution and this isn't promising
27        if best_solution and node.score < best_score * 0.7:
28            continue
29
30        # Expand
31        children = self._expand_node(node)
32        for child in children:
33            heapq.heappush(frontier, (-child.score, id(child), child))
34
35    return best_solution

1def beam_search(
2    self,
3    root: ThoughtNode,
4    beam_width: int = 5
5) -> Optional[ThoughtNode]:
6    """Beam search - keep top k nodes at each depth."""
7
8    current_beam = [root]
9    best_solution = None
10    best_score = float('-inf')
11
12    for depth in range(self.max_depth):
13        # Expand all nodes in current beam
14        all_children = []
15        for node in current_beam:
16            if self._is_solution(node.state):
17                if node.score > best_score:
18                    best_solution = node
19                    best_score = node.score
20            else:
21                children = self._expand_node(node)
22                all_children.extend(children)
23
24        if not all_children:
25            break
26
27        # Keep top beam_width nodes
28        all_children.sort(key=lambda n: n.score, reverse=True)
29        current_beam = all_children[:beam_width]
30
31        print(f"Depth {depth + 1}: {len(current_beam)} nodes, "
32              f"best score: {current_beam[0].score:.3f}")
33
34    return best_solution

1from anthropic import Anthropic
2
3class LLMEvaluator:
4    """Evaluate thoughts using LLM scoring."""
5
6    def __init__(self, model: str = "claude-sonnet-4-20250514"):
7        self.client = Anthropic()
8        self.model = model
9
10    def evaluate(
11        self,
12        problem: str,
13        thought_path: list[str],
14        current_state: dict
15    ) -> float:
16        """Evaluate how promising this thought path is."""
17
18        prompt = f"""Evaluate this partial solution to a problem.
19
20Problem: {problem}
21
22Reasoning so far:
23{chr(10).join(f"{i+1}. {t}" for i, t in enumerate(thought_path))}
24
25Current state: {current_state}
26
27Rate this partial solution:
281. Progress toward goal (0-10): How much progress has been made?
292. Correctness (0-10): Are the steps logically sound?
303. Promise (0-10): How likely is this path to lead to a good solution?
314. Efficiency (0-10): Is this an efficient approach?
32
33Return JSON:
34{{
35    "progress": X,
36    "correctness": X,
37    "promise": X,
38    "efficiency": X,
39    "overall": X,
40    "reasoning": "brief explanation"
41}}"""
42
43        response = self.client.messages.create(
44            model=self.model,
45            max_tokens=512,
46            messages=[{"role": "user", "content": prompt}]
47        )
48
49        import json
50        result = json.loads(response.content[0].text)
51        return result["overall"] / 10.0  # Normalize to 0-1

1async def evaluate_with_voting(
2    problem: str,
3    thought_path: list[str],
4    num_votes: int = 3
5) -> float:
6    """Evaluate using multiple LLM calls and voting."""
7
8    # Generate multiple evaluations in parallel
9    tasks = [
10        llm_evaluate_async(problem, thought_path)
11        for _ in range(num_votes)
12    ]
13
14    scores = await asyncio.gather(*tasks)
15
16    # Aggregate scores
17    avg_score = sum(scores) / len(scores)
18    std_dev = statistics.stdev(scores) if len(scores) > 1 else 0
19
20    # Penalize high variance (uncertain evaluations)
21    confidence_penalty = min(std_dev * 0.5, 0.2)
22
23    return max(0, avg_score - confidence_penalty)

1class MathProblemEvaluator:
2    """Heuristic evaluator for math problems."""
3
4    def evaluate(self, state: dict) -> float:
5        """Evaluate math problem state."""
6        score = 0.5  # Base score
7
8        # Reward progress toward answer
9        if "partial_result" in state:
10            score += 0.1
11
12        # Reward correct intermediate steps
13        if state.get("verified_steps", 0) > 0:
14            score += 0.1 * state["verified_steps"]
15
16        # Penalize very long solutions
17        steps = state.get("num_steps", 0)
18        if steps > 10:
19            score -= 0.1 * (steps - 10)
20
21        # Reward if close to final answer form
22        if state.get("has_numerical_result", False):
23            score += 0.2
24
25        return min(max(score, 0), 1)  # Clamp to [0, 1]
26
27
28class CodeProblemEvaluator:
29    """Heuristic evaluator for coding problems."""
30
31    def evaluate(self, state: dict) -> float:
32        """Evaluate coding problem state."""
33        score = 0.5
34
35        # Check if code is syntactically valid
36        if state.get("syntax_valid", False):
37            score += 0.15
38
39        # Check test case results
40        passed = state.get("tests_passed", 0)
41        total = state.get("tests_total", 1)
42        score += 0.3 * (passed / total)
43
44        # Penalize very long code
45        lines = state.get("code_lines", 0)
46        if lines > 50:
47            score -= 0.05 * (lines - 50) / 10
48
49        return min(max(score, 0), 1)

1class HybridEvaluator:
2    """Combine heuristic and LLM evaluation."""
3
4    def __init__(
5        self,
6        heuristic_evaluator: Callable,
7        llm_evaluator: LLMEvaluator,
8        heuristic_weight: float = 0.4
9    ):
10        self.heuristic = heuristic_evaluator
11        self.llm = llm_evaluator
12        self.heuristic_weight = heuristic_weight
13
14    async def evaluate(
15        self,
16        problem: str,
17        thought_path: list[str],
18        state: dict
19    ) -> float:
20        """Combined evaluation."""
21
22        # Fast heuristic score
23        h_score = self.heuristic(state)
24
25        # Only call LLM if heuristic is promising
26        if h_score < 0.3:
27            return h_score
28
29        # LLM evaluation
30        llm_score = await self.llm.evaluate_async(problem, thought_path, state)
31
32        # Weighted combination
33        return (
34            self.heuristic_weight * h_score +
35            (1 - self.heuristic_weight) * llm_score
36        )

1from anthropic import Anthropic
2from dataclasses import dataclass, field
3from typing import Any, Optional, Callable, List, Tuple
4import json
5import asyncio
6import heapq
7
8@dataclass
9class ToTConfig:
10    """Configuration for Tree of Thoughts."""
11    max_depth: int = 10
12    branching_factor: int = 3
13    beam_width: int = 5
14    pruning_threshold: float = 0.3
15    max_iterations: int = 100
16
17@dataclass
18class ToTNode:
19    """Node in the thought tree."""
20    id: str
21    thought: str
22    state: dict
23    score: float = 0.0
24    parent_id: Optional[str] = None
25    children_ids: list[str] = field(default_factory=list)
26    depth: int = 0
27    is_solution: bool = False
28    metadata: dict = field(default_factory=dict)
29
30    def __lt__(self, other: 'ToTNode') -> bool:
31        return self.score > other.score
32
33class TreeOfThoughtsEngine:
34    """
35    Complete Tree of Thoughts reasoning engine.
36
37    Supports multiple search strategies, LLM-based
38    thought generation and evaluation, and dynamic
39    pruning.
40    """
41
42    def __init__(
43        self,
44        config: ToTConfig = None,
45        model: str = "claude-sonnet-4-20250514"
46    ):
47        self.config = config or ToTConfig()
48        self.client = Anthropic()
49        self.model = model
50        self.nodes: dict[str, ToTNode] = {}
51        self.node_counter = 0
52
53    def _create_node(
54        self,
55        thought: str,
56        state: dict,
57        parent_id: Optional[str] = None
58    ) -> ToTNode:
59        """Create a new tree node."""
60        self.node_counter += 1
61        node_id = f"node_{self.node_counter}"
62
63        parent_depth = 0
64        if parent_id and parent_id in self.nodes:
65            parent_depth = self.nodes[parent_id].depth
66
67        node = ToTNode(
68            id=node_id,
69            thought=thought,
70            state=state,
71            parent_id=parent_id,
72            depth=parent_depth + 1
73        )
74
75        self.nodes[node_id] = node
76
77        if parent_id and parent_id in self.nodes:
78            self.nodes[parent_id].children_ids.append(node_id)
79
80        return node
81
82    def _get_path(self, node_id: str) -> list[str]:
83        """Get thought path from root to node."""
84        path = []
85        current_id = node_id
86
87        while current_id:
88            node = self.nodes.get(current_id)
89            if node:
90                path.append(node.thought)
91                current_id = node.parent_id
92            else:
93                break
94
95        return list(reversed(path))
96
97    async def generate_thoughts(
98        self,
99        problem: str,
100        current_node: ToTNode,
101        num_thoughts: int
102    ) -> list[Tuple[str, dict]]:
103        """Generate candidate next thoughts."""
104
105        path = self._get_path(current_node.id)
106
107        prompt = f"""Generate {num_thoughts} different next steps for this problem.
108
109Problem: {problem}
110
111Current reasoning path:
112{chr(10).join(f"Step {i+1}: {t}" for i, t in enumerate(path))}
113
114Current state: {json.dumps(current_node.state)}
115
116Generate {num_thoughts} DIFFERENT possible next steps.
117Each should be a distinct approach or continuation.
118
119Return JSON array:
120[
121    {{"thought": "next step description", "new_state": {{...}}}},
122    ...
123]"""
124
125        response = self.client.messages.create(
126            model=self.model,
127            max_tokens=2048,
128            messages=[{"role": "user", "content": prompt}]
129        )
130
131        try:
132            results = json.loads(response.content[0].text)
133            return [(r["thought"], r["new_state"]) for r in results]
134        except json.JSONDecodeError:
135            return []
136
137    async def evaluate_node(
138        self,
139        problem: str,
140        node: ToTNode
141    ) -> float:
142        """Evaluate how promising a node is."""
143
144        path = self._get_path(node.id)
145
146        prompt = f"""Evaluate this partial solution.
147
148Problem: {problem}
149
150Reasoning path:
151{chr(10).join(f"Step {i+1}: {t}" for i, t in enumerate(path))}
152
153Current state: {json.dumps(node.state)}
154
155Evaluate on a scale of 0-100:
156- 0-20: Dead end, unlikely to lead to solution
157- 21-40: Weak progress, may work but not promising
158- 41-60: Moderate progress, reasonable path
159- 61-80: Good progress, likely to lead to solution
160- 81-100: Excellent, very close to or is the solution
161
162Also determine if this is a COMPLETE SOLUTION.
163
164Return JSON:
165{{"score": X, "is_solution": true/false, "reasoning": "..."}}"""
166
167        response = self.client.messages.create(
168            model=self.model,
169            max_tokens=256,
170            messages=[{"role": "user", "content": prompt}]
171        )
172
173        try:
174            result = json.loads(response.content[0].text)
175            node.score = result["score"] / 100.0
176            node.is_solution = result.get("is_solution", False)
177            return node.score
178        except json.JSONDecodeError:
179            return 0.5
180
181    async def solve(
182        self,
183        problem: str,
184        initial_state: dict = None,
185        search_method: str = "beam"
186    ) -> Optional[ToTNode]:
187        """Solve problem using Tree of Thoughts."""
188
189        # Reset state
190        self.nodes = {}
191        self.node_counter = 0
192
193        # Create root
194        root = self._create_node(
195            thought=f"Starting problem: {problem}",
196            state=initial_state or {}
197        )
198        root.score = await self.evaluate_node(problem, root)
199
200        if search_method == "beam":
201            return await self._beam_search(problem, root)
202        elif search_method == "best_first":
203            return await self._best_first_search(problem, root)
204        else:
205            raise ValueError(f"Unknown search method: {search_method}")
206
207    async def _beam_search(
208        self,
209        problem: str,
210        root: ToTNode
211    ) -> Optional[ToTNode]:
212        """Beam search through thought space."""
213
214        current_beam = [root]
215        best_solution = None
216
217        for depth in range(self.config.max_depth):
218            print(f"Depth {depth + 1}, beam size: {len(current_beam)}")
219
220            next_beam = []
221
222            for node in current_beam:
223                # Check if solution
224                if node.is_solution:
225                    if best_solution is None or node.score > best_solution.score:
226                        best_solution = node
227                    continue
228
229                # Skip if below threshold
230                if node.score < self.config.pruning_threshold:
231                    continue
232
233                # Generate child thoughts
234                thoughts = await self.generate_thoughts(
235                    problem,
236                    node,
237                    self.config.branching_factor
238                )
239
240                # Create and evaluate children
241                for thought, new_state in thoughts:
242                    child = self._create_node(thought, new_state, node.id)
243                    await self.evaluate_node(problem, child)
244                    next_beam.append(child)
245
246            if not next_beam:
247                break
248
249            # Keep top beam_width nodes
250            next_beam.sort(key=lambda n: n.score, reverse=True)
251            current_beam = next_beam[:self.config.beam_width]
252
253            print(f"  Best score: {current_beam[0].score:.3f}")
254
255        return best_solution
256
257    async def _best_first_search(
258        self,
259        problem: str,
260        root: ToTNode
261    ) -> Optional[ToTNode]:
262        """Best-first search through thought space."""
263
264        frontier = [root]
265        heapq.heapify(frontier)
266        best_solution = None
267        iterations = 0
268
269        while frontier and iterations < self.config.max_iterations:
270            iterations += 1
271            node = heapq.heappop(frontier)
272
273            # Check solution
274            if node.is_solution:
275                if best_solution is None or node.score > best_solution.score:
276                    best_solution = node
277                continue
278
279            # Depth limit
280            if node.depth >= self.config.max_depth:
281                continue
282
283            # Pruning
284            if node.score < self.config.pruning_threshold:
285                continue
286
287            # Expand
288            thoughts = await self.generate_thoughts(
289                problem,
290                node,
291                self.config.branching_factor
292            )
293
294            for thought, new_state in thoughts:
295                child = self._create_node(thought, new_state, node.id)
296                await self.evaluate_node(problem, child)
297                heapq.heappush(frontier, child)
298
299        return best_solution
300
301    def get_solution_path(self, node: ToTNode) -> list[dict]:
302        """Get the full solution path."""
303        path = []
304        current_id = node.id
305
306        while current_id:
307            current_node = self.nodes.get(current_id)
308            if current_node:
309                path.append({
310                    "step": current_node.depth,
311                    "thought": current_node.thought,
312                    "state": current_node.state,
313                    "score": current_node.score
314                })
315                current_id = current_node.parent_id
316            else:
317                break
318
319        return list(reversed(path))
320
321    def visualize_tree(self, max_nodes: int = 20) -> str:
322        """Generate text visualization of the tree."""
323        lines = ["Tree of Thoughts Visualization", "=" * 40]
324
325        def format_node(node_id: str, indent: int) -> None:
326            if len(lines) > max_nodes + 2:
327                return
328
329            node = self.nodes.get(node_id)
330            if not node:
331                return
332
333            icon = "✓" if node.is_solution else "○"
334            lines.append(
335                "  " * indent +
336                f"{icon} [{node.score:.2f}] {node.thought[:50]}..."
337            )
338
339            for child_id in node.children_ids:
340                format_node(child_id, indent + 1)
341
342        # Find root
343        roots = [n for n in self.nodes.values() if n.parent_id is None]
344        for root in roots:
345            format_node(root.id, 0)
346
347        return "\n".join(lines)

1import asyncio
2
3async def main():
4    config = ToTConfig(
5        max_depth=5,
6        branching_factor=3,
7        beam_width=4,
8        pruning_threshold=0.25
9    )
10
11    engine = TreeOfThoughtsEngine(config=config)
12
13    problem = """
14    A farmer needs to cross a river with a wolf, a goat, and a cabbage.
15    The boat can only carry the farmer and one item at a time.
16    If left alone, the wolf will eat the goat, and the goat will eat the cabbage.
17    How can the farmer get everything across safely?
18    """
19
20    initial_state = {
21        "left_bank": ["farmer", "wolf", "goat", "cabbage"],
22        "right_bank": [],
23        "boat_position": "left"
24    }
25
26    print("Solving with Tree of Thoughts...")
27    solution = await engine.solve(
28        problem=problem,
29        initial_state=initial_state,
30        search_method="beam"
31    )
32
33    if solution:
34        print("\nSolution found!")
35        path = engine.get_solution_path(solution)
36        for step in path:
37            print(f"Step {step['step']}: {step['thought']}")
38            print(f"  Score: {step['score']:.2f}")
39    else:
40        print("No solution found")
41
42    print("\n" + engine.visualize_tree())
43
44asyncio.run(main())

1async def creative_brainstorm(topic: str, num_ideas: int = 10) -> list[str]:
2    """Use ToT for creative brainstorming."""
3
4    engine = TreeOfThoughtsEngine(
5        config=ToTConfig(
6            max_depth=3,
7            branching_factor=5,
8            beam_width=8
9        )
10    )
11
12    problem = f"Generate creative ideas for: {topic}"
13
14    # Use ToT to explore idea space
15    solution = await engine.solve(problem, {"ideas": []})
16
17    # Collect all unique ideas from explored nodes
18    all_ideas = set()
19    for node in engine.nodes.values():
20        if "current_idea" in node.state:
21            all_ideas.add(node.state["current_idea"])
22
23    return list(all_ideas)[:num_ideas]

1async def plan_project(goal: str, constraints: list[str]) -> list[dict]:
2    """Use ToT for project planning."""
3
4    engine = TreeOfThoughtsEngine(
5        config=ToTConfig(
6            max_depth=6,
7            branching_factor=4,
8            beam_width=5
9        )
10    )
11
12    problem = f"""
13    Goal: {goal}
14    Constraints: {constraints}
15    Create a step-by-step plan.
16    """
17
18    solution = await engine.solve(problem, {
19        "steps_planned": [],
20        "constraints_satisfied": []
21    })
22
23    if solution:
24        return engine.get_solution_path(solution)
25    return []