Pythontree_of_thought.py
Copy
1
import os
2
from dataclasses import dataclass, field
3
import openai
4
 
5
client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
6
 
7
@dataclass
8
class ThoughtNode:
9
    thought: str
10
    score: float = 0.0
11
    depth: int = 0
12
    parent: 'ThoughtNode | None' = None
13
    children: list['ThoughtNode'] = field(default_factory=list)
14
    
15
    def path(self) -> list[str]:
16
        nodes = []
17
        node = self
18
        while node:
19
            nodes.insert(0, node.thought)
20
            node = node.parent
21
        return nodes
22
 
23
def generate_thoughts(problem: str, context: str, n: int = 3) -> list[str]:
24
    """Generate n candidate next thoughts."""
25
    prompt = f"""Problem: {problem}
26
Current thinking: {context}
27
 
28
Generate {n} different ways to continue solving this problem.
29
Each thought should be a distinct approach or step.
30
Output a JSON array of {n} strings: ["thought1", "thought2", "thought3"]"""
31
    
32
    response = client.chat.completions.create(
33
        model="gpt-4o-mini",
34
        messages=[{"role": "user", "content": prompt}],
35
        max_tokens=300, temperature=0.8,
36
        response_format={"type": "json_object"},
37
    )
38
    import json
39
    content = json.loads(response.choices[0].message.content)
40
    thoughts = list(content.values())[0] if isinstance(list(content.values())[0], list) else list(content.values())
41
    return thoughts[:n]
42
 
43
def score_thought(problem: str, thought_path: list[str]) -> float:
44
    """Score how promising a thought path is (1-10)."""
45
    path_str = " → ".join(thought_path)
46
    prompt = f"""Problem: {problem}
47
Thought path: {path_str}
48
 
49
Score this reasoning path 1-10. 10 = definitely on track to solve it.
50
Output JSON: {{"score": <1-10>, "reasoning": "..."}}"""
51
    
52
    import json
53
    response = client.chat.completions.create(
54
        model="gpt-4o-mini",
55
        messages=[{"role": "user", "content": prompt}],
56
        max_tokens=100, temperature=0,
57
        response_format={"type": "json_object"},
58
    )
59
    result = json.loads(response.choices[0].message.content)
60
    return float(result.get("score", 1))
61
 
62
def bfs_tot(problem: str, depth: int = 3, breadth: int = 3, prune_threshold: float = 5.0) -> ThoughtNode | None:
63
    """BFS Tree of Thought: explore width-first, prune low-scoring branches."""
64
    root = ThoughtNode(thought=f"Problem: {problem}", depth=0)
65
    frontier = [root]
66
    best_node = None
67
    
68
    for d in range(depth):
69
        next_frontier = []
70
        
71
        for node in frontier:
72
            context = " → ".join(node.path())
73
            thoughts = generate_thoughts(problem, context, n=breadth)
74
            
75
            for thought in thoughts:
76
                child = ThoughtNode(thought=thought, depth=d+1, parent=node)
77
                child.score = score_thought(problem, child.path())
78
                node.children.append(child)
79
                
80
                if child.score >= prune_threshold:
81
                    next_frontier.append(child)
82
                    if best_node is None or child.score > best_node.score:
83
                        best_node = child
84
                        
85
                print(f"  Depth {d+1}: [{child.score:.0f}/10] {thought[:60]}...")
86
        
87
        frontier = sorted(next_frontier, key=lambda n: n.score, reverse=True)[:breadth]
88
        if not frontier:
89
            print("  All branches pruned.")
90
            break
91
    
92
    return best_node
93
 
94
# ── Run ToT ──────────────────────────────────────────────────────────
95
PROBLEM = "Design a system to fairly allocate limited hospital ICU beds during a crisis with 3x demand."
96
 
97
print(f"=== Tree of Thought ===")
98
print(f"Problem: {PROBLEM}\n")
99
best = bfs_tot(PROBLEM, depth=2, breadth=2)
100
 
101
if best:
102
    print(f"\n=== Best Path (score={best.score:.0f}/10) ===")
103
    for i, thought in enumerate(best.path()):
104
        print(f"Step {i}: {thought[:100]}")
105