Implementing the Tree of Thoughts in LangChain’s Chain

Introduction

Tree of Thoughts (ToT) is an algorithm that combines Large Language Models (LLMs) and heuristic search, as presented in this paper by Princeton University and Google DeepMind. It appears that this algorithm is being implemented into Gemini, a multimodal generative AI that is currently under development by Google.

Now, let’s briefly touch on the thought process of ToT (Tree of Thoughts).

In the usual CoT (Chain of Thoughts) approach, LLMs tend to progress linearly in their thinking towards problem solving, and if an error occurs along the way, they tend to proceed along that erroneous criterion.

In contrast, in the ToT (Tree of Thoughts) approach, LLMs evaluate themselves at each stage of thought and stop inefficient approaches early, switching to alternative methods.

Implementation

In implementing this in LangChain’s Chain, I utilized the prompt provided by the poster of the following video. To explain the process simply, it first generates a wide range of ideas, then evaluates each idea, delves into it, and selects the one with the most promise of success.

PromptTemplate is used to define the Tree of Thoughts prompt, and the chain is implemented at each step.

from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI


template ="""
Step1 :
 
I have a problem related to {input}. Could you brainstorm three distinct solutions? Please consider a variety of factors such as {perfect_factors}
A:
"""

prompt = PromptTemplate(
    input_variables=["input","perfect_factors"],
    template = template                      
)

chain1 = LLMChain(
    llm=ChatOpenAI(temperature=0, model="gpt-4"),
    prompt=prompt,
    output_key="solutions"
)

template ="""
Step 2:

For each of the three proposed solutions, evaluate their potential. Consider their pros and cons, initial effort needed, implementation difficulty, potential challenges, and the expected outcomes. Assign a probability of success and a confidence level to each option based on these factors

{solutions}

A:"""

prompt = PromptTemplate(
    input_variables=["solutions"],
    template = template                      
)

chain2 = LLMChain(
    llm=ChatOpenAI(temperature=0, model="gpt-4"),
    prompt=prompt,
    output_key="review"
)

template ="""
Step 3:

For each solution, deepen the thought process. Generate potential scenarios, strategies for implementation, any necessary partnerships or resources, and how potential obstacles might be overcome. Also, consider any potential unexpected outcomes and how they might be handled.

{review}

A:"""

prompt = PromptTemplate(
    input_variables=["review"],
    template = template                      
)

chain3 = LLMChain(
    llm=ChatOpenAI(temperature=0, model="gpt-4"),
    prompt=prompt,
    output_key="deepen_thought_process"
)

template ="""
Step 4:

Based on the evaluations and scenarios, rank the solutions in order of promise. Provide a justification for each ranking and offer any final thoughts or considerations for each solution
{deepen_thought_process}

A:"""

prompt = PromptTemplate(
    input_variables=["deepen_thought_process"],
    template = template                      
)

chain4 = LLMChain(
    llm=ChatOpenAI(temperature=0, model="gpt-4"),
    prompt=prompt,
    output_key="ranked_solutions"
)

We connect the four chains using ‘SequentialChain’. The output of one chain becomes the input to the next chain.

from langchain.chains import SequentialChain

overall_chain = SequentialChain(
    chains=[chain1, chain2, chain3, chain4],
    input_variables=["input", "perfect_factors"],
    output_variables=["ranked_solutions"],
    verbose=True
)

print(overall_chain({"input":"human colonization of Mars", "perfect_factors":"The distance between Earth and Mars is very large, making regular resupply difficult"}))

Output:

{
    "input": "human colonization of Mars",
    "perfect_factors": "The distance between Earth and Mars is very large, making regular resupply difficult",
    "ranked_solutions": {
        "Ranking_1": {
            "Justification": "Using In-Situ Resource Utilization is the most promising solution due to its potential to provide the necessary resources for a Mars colony and reduce the need for resupply missions from Earth. The medium initial effort, implementation difficulty, and potential challenges are outweighed by the high probability of success and 70% confidence level.",
            "In_Situ_Resource_Utilization_ISRU": {
                "Pros": "This solution could provide the necessary resources for a Mars colony and reduce the need for resupply missions from Earth.",
                "Cons": "ISRU is technically challenging and would require significant investment in research and development.",
                "Initial_Effort": "Medium. This would require the development of new technology and the establishment of infrastructure on Mars.",
                "Implementation_Difficulty": "Medium. ISRU is a complex task that requires advanced technology.",
                "Potential_Challenges": "Technical difficulties, high costs.",
                "Expected_Outcomes": "If successful, ISRU could provide a steady supply of resources for a Mars colony.",
                "Probability_of_Success": "High. ISRU is already being tested by NASA and other space agencies.",
                "Confidence_Level": "70%"
            }
        },
        "Ranking_2": {
            "Justification": "Building a self-sustaining colony is a promising solution due to its potential to make the Mars colony self-sufficient. However, the high initial effort, implementation difficulty, and potential challenges make it less promising than the first solution. The medium probability of success and 60% confidence level also contribute to its ranking.",
            "Building_a_Self_Sustaining_Colony": {
                "Pros": "This solution could make the Mars colony self-sufficient, reducing the need for resupply missions from Earth.",
                "Cons": "Building a self-sustaining colony is a complex task that requires advanced technology and a lot of resources.",
                "Initial_Effort": "High. This would require the development of new technology and the establishment of infrastructure on Mars.",
                "Implementation_Difficulty": "High. Building a self-sustaining colony is a complex task that requires advanced technology.",
                "Potential_Challenges": "Technical difficulties, high costs.",
                "Expected_Outcomes": "If successful, a self-sustaining colony could reduce the need for resupply missions from Earth.",
                "Probability_of_Success": "Medium. While there are significant challenges, there is also a lot of interest in building a self-sustaining colony on Mars.",
                "Confidence_Level": "60%"
            }
        },
        "Ranking_3": {
            "Justification": "While asteroid mining has the potential to provide a steady supply of resources for a Mars colony, the high initial effort, implementation difficulty, and potential challenges make it a less promising solution compared to others. The medium probability of success and 50% confidence level also contribute to its lower ranking.",
            "Terraforming_Mars": {
                "Pros": "This solution could make Mars more habitable for humans, reducing the need for life support systems and making the colony more self-sufficient.",
                "Cons": "Terraforming is a long-term process that could take centuries or even millennia. It would also require a massive amount of resources and energy.",
                "Initial_Effort": "Extremely High. Terraforming would require a massive amount of resources and energy.",
                "Implementation_Difficulty": "Extremely High. Terraforming is a long-term process that could take centuries or even millennia.",
                "Potential_Challenges": "Technical difficulties, high costs, time scale.",
                "Expected_Outcomes": "If successful, terraforming could make Mars more habitable for humans.",
                "Probability_of_Success": "Low. Terraforming is a theoretical concept and has never been attempted before.",
                "Confidence_Level": "20%"
            }
        }
    }
}

Finally, I thought that not only using the tree of thoughts in the prompt, but also introducing it into the planning process of the Agent could improve the accuracy of tool selection and planning.