Prompt EngineeringAdvanced

Implement a ToT Prompt Pipeline

Propose several next steps, judge the partial states, keep the best and expand them. How to build a Tree-of-Thought search, why the evaluator decides everything, and when a reasoning model makes it unnecessary.

By Amit UpadhyayPublished 25 Sep 202610 min

Following one scenario: “Build next week’s on-call rota: nobody twice in a row, Priya off Wednesday, at least one senior every night.”An on-call scheduler for a six-person team with hard constraintsWhy it matters →

What you will build

A tree-of-thought rota builder with code-checked constraints, bounded search and a solver comparison.

The same scenario as Read and See. You are done when every line below is true. Tick them off as you go.

Done when

The scenario week produces a valid rota, with every constraint checked in code
Invalid branches are pruned as soon as they break a constraint
Calls per search are logged and bounded
Results are compared with a single-chain prompt and a constraint solver

0 of 4 done

0/1 steps

Build Steps

Progress

0/1

◆ Before you start

Prerequisites

Python 3.11+Required
python --version
openaiRequired
pip install openai

Step 1 of 1

Implement Tree of Thought reasoning

Build a Tree of Thought (ToT) solver that explores multiple reasoning paths and selects the best.

Tasks

Implement BFS-ToT that generates 3 thoughts per step

Write a thought evaluator that scores 1-10 (viable/promising)

Prune low-scoring branches (score < 5)

Compare ToT vs CoT on a logic puzzle

Visualize the thought tree with depth and scores

🎯 What You'll Learn

▸Generate multiple thought branches for a problem
▸Score thought paths with a value function
▸Use BFS/DFS to find the highest-value solution path

Code

Pythontree_of_thought.py
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-6-luna", reasoning_effort="none",
34         messages=[{"role": "user", "content": prompt}],
35         max_completion_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-6-luna", reasoning_effort="none",
55         messages=[{"role": "user", "content": prompt}],
56         max_completion_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

💡 Pro Tips

▸ToT is significantly more expensive than CoT (N×M model calls) — use for hard problems only
▸The key insight: thinking with multiple perspectives improves solution quality for creative/planning tasks

Evaluate it

Prove it works before you call it done

Working once on your machine is a demo. Measure it against a fixed set of cases.

Generate 20 weeks of constraints, including 2 that cannot be satisfied. Run single-chain, tree-of-thought and a solver.

Pass bar

Tree-of-thought finds valid rotas for all 18 feasible weeks and reports both infeasible ones clearly

Ship and observe

What changes when it is live

Before it ships, in order

Report conflicting constraints instead of returning nothing
Keep the solver as a fallback for fixed-rule weeks

Put on a dashboard

Valid first-search rate
Calls per search
Violations reaching people

Break it on purpose

Learn the failure modes while they are cheap

Make one change at a time, predict what will happen, then run it.

ChangeWhat you should see

Let the model judge branch validityInvalid rotas start getting through.

Remove the depth limit and use contradictory constraintsCall counts explode before the search gives up.

Use a single chain insteadViolations appear late in the week.

What this build costs

Twenty test weeks cost a fraction of a cent.

Where this leads

Agent design patterns

Search over thoughts is one strategy. Agent loops add actions and real-world feedback to the mix.

Next topic →

tree-of-thought ToT reasoning search beam search prompt-engineering

← PreviousConstrained Generation: Beyond Hope-Based Parsing Next →Red Teaming LLM Applications: Attack It Before Your Users Do

Continue reading

All Prompt Engineering articles →

Prompt Engineering9 min

Meta-Prompting: Letting a Model Tune Your Prompts

From a one-off rewrite to refinement loops on failures, candidate search and compiled DSPy programs. How to optimise prompts against a metric without overfitting, and how to review what the optimiser produces.

Prompt Engineering9 min

Red Teaming LLM Applications: Attack It Before Your Users Do

Threat model first, then manual sessions, then automated scans with garak, promptfoo and PyRIT. How to find jailbreaks, data leaks and unsafe tool calls, fix them in the architecture, and keep them fixed in CI.

Prompt Engineering13 min

Constrained Generation: Beyond Hope-Based Parsing

Enforce output structure at the token level with grammars and token masks — how it differs from native JSON mode, and when the strictness costs you reasoning quality.

Discussion

A tree-of-thought rota builder with code-checked constraints, bounded search and a solver comparison.

The same scenario as Read and See. You are done when every line below is true. Tick them off as you go.

Done when

The scenario week produces a valid rota, with every constraint checked in code
Invalid branches are pruned as soon as they break a constraint
Calls per search are logged and bounded
Results are compared with a single-chain prompt and a constraint solver

0 of 4 done

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-6-luna", reasoning_effort="none", 34 messages=[{"role": "user", "content": prompt}], 35 max_completion_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-6-luna", reasoning_effort="none", 55 messages=[{"role": "user", "content": prompt}], 56 max_completion_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

Prove it works before you call it done

Working once on your machine is a demo. Measure it against a fixed set of cases.

Generate 20 weeks of constraints, including 2 that cannot be satisfied. Run single-chain, tree-of-thought and a solver.

Pass bar

Tree-of-thought finds valid rotas for all 18 feasible weeks and reports both infeasible ones clearly

Learn the failure modes while they are cheap

Make one change at a time, predict what will happen, then run it.

ChangeWhat you should see

Let the model judge branch validityInvalid rotas start getting through.

Remove the depth limit and use contradictory constraintsCall counts explode before the search gives up.

Use a single chain insteadViolations appear late in the week.

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-6-luna", reasoning_effort="none",
34	messages=[{"role": "user", "content": prompt}],
35	max_completion_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-6-luna", reasoning_effort="none",
55	messages=[{"role": "user", "content": prompt}],
56	max_completion_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