How to Implement a Tree of Thoughts in Python

The concept of the "Tree of Thoughts" (ToT) is a powerful tool for allowing an LLM to explore a non-linear path through a problem space, particularly for exploring and evolving ideas systematically. This blog post tutorial walks you through the details of implementing a Tree of Thoughts (ToT) prompting technique in Python, leveraging Anthropic's Claude Sonnet 3.5 language model to generate and expand thoughts. The following example provides a practical illustration of how to build and utilize a ToT to ideate solutions efficiently.

All code for this tutorial can be found in my GitHub repository.

What is a Tree of Thoughts?

A Tree of Thoughts is a hierarchical structure where each node represents a distinct thought or idea. The tree grows by expanding nodes with additional thoughts, creating branches that represent different avenues of exploration. This approach allows for a structured and systematic exploration of ideas, helping to uncover innovative solutions and insights.

Implementation Overview

We'll walk through a Python implementation of a Tree of Thoughts, explaining each component and how it interacts with the Anthropic Claude Sonnet 3.5 API to generate and evolve thoughts.

Setting Up the Environment

First, ensure you have the necessary libraries installed. You'll need anthropic and python-dotenv. Install these using pip if you haven't already:

pip install anthropic python-dotenv

Importing Libraries and Setting Up Anthropic

Ensure you have a .env file with your Anthropic API key, such as the following:

ANTHROPIC_API_KEY=YOUR_ANTHROPIC_API_KEY

The AnthropicService class that uses your Anthropic API key is a wrapper around the Anthropic API that provides a simple interface for interacting with the API. In our ai.py file, we define the AnthropicService class:

import os
from dotenv import load_dotenv
from anthropic import Anthropic

# Load environment variables from .env file
load_dotenv()

MAX_TOKENS = 1024
DEFAULT_TEMPERATURE = 0
DEFAULT_ANTHROPIC_MODEL_NAME = "claude-3-5-sonnet-20240620"


class AnthropicService:
    def __init__(self, model_name: str = None, anthropic: Anthropic = None):
        if not os.environ.get("ANTHROPIC_API_KEY"):
            raise ValueError("No valid API key found for Anthropic.")
        self.client = anthropic or Anthropic(
            api_key=os.environ["ANTHROPIC_API_KEY"])
        self.model_name = model_name or DEFAULT_ANTHROPIC_MODEL_NAME

    def generate_response(self, prompt: str, max_tokens: int = MAX_TOKENS, temperature: float = DEFAULT_TEMPERATURE) -> str:
        msg = self.client.messages.create(
            model=self.model_name,
            max_tokens=max_tokens,
            temperature=temperature,
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                },
            ],
        )
        return msg.content[0].text

Import the AnthropicService class into app.py to use it:

from ai import AnthropicService

Defining the ThoughtNode Class

The ThoughtNode class represents a node in the Tree of Thoughts. Each node contains a thought and a list of child nodes.

class ThoughtNode:
    def __init__(self, thought, children=None):
        self.thought = thought
        self.children = children or []

Building the TreeOfThought Class

The TreeOfThought class orchestrates the process of generating and evolving thoughts. It initializes with a root prompt and interacts with the AI service to expand the tree iteratively.

class TreeOfThought:
    def __init__(self, root_prompt, ai_service=None, max_iterations=3, max_tokens=250):
        self.root = ThoughtNode(root_prompt)
        self.max_iterations = max_iterations
        self.ai_service = ai_service or AnthropicService()
        self.current_thoughts = [self.root]
        self.max_tokens = max_tokens

Calling the Language Model

The call_llm method sends a prompt to the AI service and returns the response. This method handles any errors that may occur during the API call.

    def call_llm(self, prompt):
        try:
            response = self.ai_service.generate_response(
                prompt,
                max_tokens=self.max_tokens,
            )
            return response
        except Exception as e:
            print(f"Error calling LLM: {e}")
            return []

Exploring and Expanding Thoughts

The explore_thoughts method generates new thoughts based on the current thoughts in the tree. It prompts the AI to provide two next thoughts for each current thought, creating new nodes and expanding the tree.

    def explore_thoughts(self, thought_nodes):
        new_thought_nodes = []
        for thought_node in thought_nodes:
            prompt = f"Given the current thought: '{thought_node.thought}', provide two concise next thoughts that evolve this idea further."
            response = self.call_llm(prompt)
            if response:
                new_thought_node = ThoughtNode(response)
                thought_node.children.append(new_thought_node)
                new_thought_nodes.append(new_thought_node)
        return new_thought_nodes

Running the Tree of Thoughts

The run method orchestrates the iterative process of expanding the tree. It continues exploring thoughts until the maximum number of iterations is reached.

    def run(self):
        iteration = 0
        while self.current_thoughts and iteration < self.max_iterations:
            print(f"Iteration {iteration + 1}:")
            self.current_thoughts = self.explore_thoughts(
                self.current_thoughts)
            for thought_node in self.current_thoughts:
                print(f"Explored Thought: {thought_node.thought}")
            iteration += 1

Updating and Printing the Tree of Thoughts

The update_starting_thought method allows for changing the root thought of the tree. The print_tree method provides a visual representation of the entire tree, showing the hierarchy of thoughts.

    def update_starting_thought(self, new_thought):
        self.root = ThoughtNode(new_thought)
        self.current_thoughts = [self.root]

    def print_tree(self, node, level=0):
        indent = ' ' * (level * 2)
        thought_lines = node.thought.split('\n')
        for idx, line in enumerate(thought_lines):
            if idx == 0:
                print(f"{indent}- {line}")
            else:
                print(f"{indent}  {line}")
        for child in node.children:
            self.print_tree(child, level + 1)

Execution

Finally, we set the starting prompt and run the Tree of Thoughts, printing the final structure.

if __name__ == "__main__":
    starting_prompt = "Think of a solution to reduce the operational costs of your business."
    tot = TreeOfThought(starting_prompt)
    tot.run()
    print("=" * 100)
    print("Final Tree of Thoughts:")
    tot.print_tree(tot.root)

You should see output of the tree of thoughts like this:

Final Tree of Thoughts:
- Think of a solution to reduce the operational costs of your business.
  - Here are two concise next thoughts that evolve the idea of reducing operational costs:

    1. Analyze energy consumption and implement efficiency measures.

    2. Explore automation options for repetitive tasks to reduce labor costs.
    - Here are two concise next thoughts that further evolve the idea of reducing operational costs:

      1. Implement a lean inventory management system to minimize holding costs and waste.

      2. Negotiate better terms with suppliers and consider strategic partnerships for bulk purchasing discounts.
      - Here are two concise next thoughts that further evolve the idea of reducing operational costs:

        1. Adopt energy-efficient technologies and practices to lower utility expenses and reduce environmental impact.

        2. Invest in automation and AI-driven processes to streamline operations and reduce labor costs in the long term.

Conclusion

This Python implementation of a Tree of Thoughts prompting technique demonstrates how to systematically explore and evolve ideas using Anthropic's Claude Sonnet 3.5 language model. By structuring thoughts in a hierarchical tree, you can uncover innovative solutions and insights efficiently. This approach is particularly valuable for brainstorming, problem-solving, and any scenario where exploring multiple avenues of thought is beneficial.

Whether you're working on business strategies, creative writing, research, or any other problem involving non-linear thinking, the Tree of Thoughts (ToT) prompting technique offers a structured and powerful method for ideation and exploration.