Tree of Thoughts: Deliberate Problem Solving with Large Language Models¶

# Generate 3 different next steps
thought_1 = llm.generate("Next step: ...")  # Independent sample
thought_2 = llm.generate("Next step: ...")  # Independent sample
thought_3 = llm.generate("Next step: ...")  # Independent sample

# Ask for 3 alternatives at once
thoughts = llm.generate("""
Propose 3 different next steps:
1.
2.
3.
""")

for thought in candidate_thoughts:
    score = llm.evaluate(f"""
    On a scale of 1-10, how likely is this path to lead to the solution?
    Current state: {current_state}
    Next thought: {thought}
    Score (1-10):
    """)

prompt = f"""
Which of these paths is most promising?
A) {thought_a}
B) {thought_b}
C) {thought_c}
Vote (A/B/C):
"""
votes = [llm.vote(prompt) for _ in range(5)]  # Get multiple votes
best = most_common(votes)

def tree_of_thoughts_bfs(problem, b=3, max_depth=5):
    """
    Tree of Thoughts with Breadth-First Search.

    Args:
        problem: Initial problem state
        b: Branching factor (thoughts per state)
        max_depth: Maximum search depth
    """
    # Initialize with root
    current_level = [problem]

    for depth in range(max_depth):
        next_level = []

        for state in current_level:
            # Check if this state is a solution
            if is_solution(state):
                return state

            # Generate b candidate next thoughts
            thoughts = generate_thoughts(state, num_thoughts=b)

            # Evaluate each thought
            values = [evaluate_state(state, thought) for thought in thoughts]

            # Create new states
            for thought, value in zip(thoughts, values):
                new_state = apply_thought(state, thought)
                next_level.append((new_state, value))

        # Keep top-k most promising states
        next_level.sort(key=lambda x: x[1], reverse=True)
        current_level = [s for s, v in next_level[:b]]

        if not current_level:
            return None  # No solution found

    return None  # Max depth reached


def generate_thoughts(state, num_thoughts):
    """Generate multiple candidate thoughts."""
    prompt = f"""
    Current state: {state}
    Propose {num_thoughts} different next steps:
    """
    response = llm.generate(prompt)
    return parse_thoughts(response)


def evaluate_state(state, thought):
    """Evaluate how promising a state is."""
    prompt = f"""
    State: {state}
    Next thought: {thought}

    Rate the likelihood this leads to solution (1-10):
    """
    score = llm.generate(prompt)
    return int(score)

def tree_of_thoughts_dfs(state, depth=0, max_depth=5):
    """
    Tree of Thoughts with Depth-First Search.
    """
    # Base cases
    if is_solution(state):
        return state
    if depth >= max_depth:
        return None

    # Generate and evaluate thoughts
    thoughts = generate_thoughts(state, num_thoughts=3)
    values = [evaluate_state(state, t) for t in thoughts]

    # Sort by value (best first)
    sorted_thoughts = sorted(zip(thoughts, values),
                            key=lambda x: x[1],
                            reverse=True)

    # Try each thought (DFS)
    for thought, value in sorted_thoughts:
        # Prune obviously bad paths
        if value < threshold:
            continue

        new_state = apply_thought(state, thought)
        result = tree_of_thoughts_dfs(new_state, depth+1, max_depth)

        if result is not None:
            return result  # Found solution

    return None  # No solution in this subtree

# ToT can explore different solution strategies
problem = "Prove that sqrt(2) is irrational"

# Strategy 1: Proof by contradiction
# Strategy 2: Use fundamental theorem of arithmetic
# Strategy 3: Continued fractions

# Evaluate which approach is most promising
# Pursue best strategy
# Backtrack if stuck

from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate

class TreeOfThoughts:
    def __init__(self, llm, branching_factor=3):
        self.llm = llm
        self.b = branching_factor

    def solve(self, problem, max_depth=5):
        """Solve problem using Tree of Thoughts."""
        return self._bfs(problem, max_depth)

    def _bfs(self, initial_state, max_depth):
        """Breadth-first search."""
        states = [{"state": initial_state, "history": []}]

        for depth in range(max_depth):
            next_states = []

            for current in states:
                # Check if solved
                if self._is_solution(current["state"]):
                    return current

                # Generate thoughts
                thoughts = self._generate_thoughts(current["state"])

                # Evaluate and expand
                for thought in thoughts:
                    value = self._evaluate(current["state"], thought)

                    if value > 5:  # Threshold
                        new_state = self._apply_thought(
                            current["state"], thought
                        )
                        next_states.append({
                            "state": new_state,
                            "history": current["history"] + [thought],
                            "value": value
                        })

            # Keep top-k
            next_states.sort(key=lambda x: x["value"], reverse=True)
            states = next_states[:self.b]

            if not states:
                return None

        return None

    def _generate_thoughts(self, state):
        """Generate candidate next thoughts."""
        prompt = f"""
        Current state: {state}

        Propose {self.b} different next steps to solve this problem.
        Each step should be distinct and promising.

        Steps:
        1.
        2.
        3.
        """
        response = self.llm(prompt)
        return self._parse_thoughts(response)

    def _evaluate(self, state, thought):
        """Evaluate thought value."""
        prompt = f"""
        Current state: {state}
        Proposed next step: {thought}

        On a scale of 1-10, how likely is this step to lead to the solution?
        Consider:
        - Does it make progress toward the goal?
        - Is it a valid step?
        - Does it avoid dead ends?

        Rating (1-10):
        """
        response = self.llm(prompt)
        try:
            return int(response.strip())
        except:
            return 5  # Default

    def _parse_thoughts(self, response):
        """Parse thoughts from LLM response."""
        lines = response.strip().split('\n')
        thoughts = []
        for line in lines:
            if line.strip() and line[0].isdigit():
                # Remove numbering
                thought = line.split('.', 1)[1].strip()
                thoughts.append(thought)
        return thoughts[:self.b]

    def _apply_thought(self, state, thought):
        """Apply thought to state."""
        return f"{state}\n{thought}"

    def _is_solution(self, state):
        """Check if state is a solution."""
        prompt = f"""
        State: {state}

        Is this a complete solution to the problem? (yes/no)
        """
        response = self.llm(prompt).strip().lower()
        return 'yes' in response


# Usage
llm = OpenAI(temperature=0.7, model="gpt-4")
tot = TreeOfThoughts(llm, branching_factor=3)

problem = """
Use the numbers 4, 9, 10, 13 exactly once each with operations +, -, ×, ÷ to make 24.
"""

solution = tot.solve(problem, max_depth=3)
if solution:
    print("Solution found!")
    print("\n".join(solution["history"]))
else:
    print("No solution found")

# Don't use full ToT for everything
# Start with CoT, upgrade if needed

if is_simple_problem(task):
    return chain_of_thought(task)
elif needs_search(task):
    return tree_of_thoughts(task)
else:
    return direct_prompting(task)

# Higher b = better quality, higher cost
# Lower b = faster, may miss solutions

easy_tasks: b = 2
medium_tasks: b = 3
hard_tasks: b = 5

# Don't keep bad paths
if evaluate(thought) < threshold:
    continue  # Skip this branch

# Prune percentage at each level
keep_top_percent = 0.5  # Keep top 50%

# Don't re-evaluate same states
evaluation_cache = {}

def evaluate_cached(state, thought):
    key = hash(state + thought)
    if key not in evaluation_cache:
        evaluation_cache[key] = evaluate(state, thought)
    return evaluation_cache[key]