How to Implement a Tree of Thoughts in Python
Tags: treeofthought, ai, python, anthropic
Published: 2024-07-05 20:36:00

---
title: How to Implement a Tree of Thoughts in Python
published: true
description: This blog post walks you through the implementation of a Tree of Thoughts (ToT) prompting technique in Python, a powerful tool for systematically exploring and evolving ideas. Using Anthropic's Claude Sonnet 3.5 language model, we build a hierarchical structure where each node represents a thought, and the tree expands by generating new thoughts. Whether you're looking to brainstorm, problem-solve, or innovate, this post guides you through setting up and utilizing a Tree of Thoughts in your projects.
tags:
- TreeofThought
- AI
- Python
- Anthropic
# cover_image: https://direct_url_to_image.jpg
# Use a ratio of 100:42 for best results.
published_at: 2024-07-05 20:36 +0000
---

![Robot sitting on a tree limb](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/f58uiniwpolxysezd5ef.jpg)

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](https://github.com/stephenc222/example-tree-of-thoughts-prompting).

## 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:

```sh
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:

```bash
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:

```python
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:

```python
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.

```python
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.

```python
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.

```python
    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.

```python
    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.

```python
    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.

```python
    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.

```python
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:

```txt
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.