Implementing Tree-of-Thoughts(ToT) with LLM using Langchain
In a landscape where language models often seem to possess a mind of their own, a new frontier emerges with the Tree of Thoughts approach, introduced by Google DeepMind and Princeton University. Beyond mere answering, this model delve into self-evaluation, navigating a labyrinth of choices to determine the optimal path forward at each step of problem solving. Much like a seasoned strategist, they engage in foresight and retrospection, unraveling a web of possibilities to provide solutions with unprecedented depth and clarity. Welcome to the realm where language models not only speak but also contemplate, anticipate, and strategize.
How ToT works?
- Given a problem that has non-trivial solutions and planning (eg: Game of 24, Creative Writing, Mini Crosswords), ToT explores various reasoning paths independently.
- Assess each path and choices autonomously to determine the next course of action.
- Intermediate results presented in high-level human-readable language.
- Evaluate multiple potential feasible plans to solve problems.
- Proceed with the most viable plan
Chain of Thoughts prompting helps to solve problems linearly, however if it makes error at any of the steps, it gets carried over to next steps until the end of execution. Whereas, Tree of Thoughts technique aims to overcome this limitation by self-evaluating solutions at each step and eliminates the path that doesn’t work out.
Implementation of ToT using Langchain
Let’s dive into how we can implement a basic ToT in Python using Langchain. Langchain helps to build and deploy LLM and provides support to use almost any models like ChatGPT, Claude, etc. I have chosen a creative writing task to plan and evaluate air taxi implementation using ToT.
Problem Description
For problem solving, it’s vital to rationalize solutions considering various factors and weighing the pros and cons of each option. I am going to construct an LLM for each of these critical thinking step —
Step 1: Provide solutions for implementing air taxi in urban areas
Step 2: Analyse the advantages and disadvantages, implementation and planning strategies for each of the suggested solutions
Step 3: Evaluate the proposed plans
RunnableSequence is used to construct LLM chain. Using Runnable Sequence, we can pass the output of one LLM as input to the next one. It can be used by calling a ‘|’ operator.
from langchain.llms import OpenAI
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.schema.runnable import RunnablePassthrough
# Step 1 - Building LLM to draft solutions
solutions_template = """
Step 1:
I'm facing a challenge related to {topic}. Can you generate two plans for implementing {topic}?
"""
solutions_prompt = PromptTemplate(
input_variables = ["topic"],
template = solutions_template
)
solutions_llm = ChatOpenAI(openai_api_key = os.getenv('OPENAI_API_KEY'), temperature = 0.1)
solutions_chain = solutions_prompt | solutions_llm
# Step 2 - Build LLM to analyse each solution drafted in the previous step
analysis_template = """
Step 2:
For each of the {solutions}, consider the advantages and disadvantages, ways to integrate them with existing urban infrastructure and planning strategies.
Your output should be in the form of description of the plan, advantages, disadvantages and planning strategies in a JSON format for each solution.
"""
analysis_prompt = PromptTemplate(input_variables = ["solutions"],
template = analysis_template)
analysis_llm = ChatOpenAI(openai_api_key = os.getenv('OPENAI_API_KEY'), temperature = 0.1)
analysis_chain = analysis_prompt | analysis_llm
# Step 3 - Build LLM to evaluate and rank solutions
evaluation_template = """
Step 3:
Evaluate and rank solutions based on {solutions} and {analysis} in previous steps,
rank the plans on a scale of 10 based on safety, environmental impact, reliability and cost-effectiveness and justify the rank.
Your output should be in the form of description about the plan, advantages, disadvantages, planning strategies,
evaluation with reason in a JSON format for each solution.
"""
evaluation_prompt = PromptTemplate(input_variables = ["solutions", "analysis"],
template = evaluation_template)
evaluation_llm = ChatOpenAI(openai_api_key = os.getenv('OPENAI_API_KEY'), temperature = 0.1)
evaluation_chain = evaluation_prompt | evaluation_llmNow, it’s time to call all the individual chains in the previous step and build a sequential chain where output of the first one is passed as an input to the next.
complete_chain = RunnablePassthrough.assign(solutions = solutions_chain)| RunnablePassthrough.assign(analysis = analysis_chain) | evaluation_chain
response = complete_chain.invoke({"topic": "air taxi"})
print(response.content)The output will be something like below —
{
"Plan 1": {
"description": "Plan 1 focuses on research and development, partnership and collaboration, marketing and promotion, operational setup, and expansion and growth of air taxi services.",
"advantages": [
"Thorough research ensures safety and efficiency of air taxi vehicles",
"Partnerships leverage expertise and resources from various industries",
"Strong marketing strategy raises awareness and attracts customers",
"Established operational hubs and safety protocols ensure smooth operations",
"Continuous improvement and expansion cater to growing market demand"
],
"disadvantages": [
"High initial investment required for research, development, and operational setup",
"Dependence on partnerships may lead to conflicts or delays",
"Marketing costs can be high and may not guarantee immediate returns",
"Expansion may be limited by regulatory approvals and infrastructure constraints",
"Competition from other transportation services may impact growth"
],
"planning_strategies": [
"Integrate operational hubs with existing transportation networks",
"Collaborate with city planners to secure necessary permits for infrastructure development",
"Utilize customer feedback to improve services and tailor expansion plans",
"Consider partnerships with existing businesses to enhance customer experience",
"Implement dynamic pricing models to stay competitive in the market"
],
"evaluation": {
"safety": 8,
"environmental impact": 7,
"reliability": 9,
"cost-effectiveness": 6
},
"reason": "Plan 1 scores high on safety and reliability due to thorough research, established safety protocols, and continuous improvement. However, it may have a moderate environmental impact and lower cost-effectiveness due to high initial investment and marketing costs."
},
"Plan 2": {
"description": "Plan 2 focuses on infrastructure development, fleet acquisition, regulatory compliance, pricing strategy, and customer experience for air taxi services.",
"advantages": [
"Dedicated infrastructure improves accessibility and efficiency of air taxi services",
"Electric eVTOL aircraft are environmentally friendly and efficient",
"Regulatory compliance ensures safety and reliability of services",
"Competitive pricing strategy attracts customers and ensures profitability",
"Focus on customer experience enhances satisfaction and loyalty"
],
"disadvantages": [
"High upfront costs for infrastructure development and fleet acquisition",
"Regulatory hurdles may delay launch or expansion of services",
"Pricing strategy may not always align with market demand",
"Customer experience improvements may require ongoing investment",
"Limited flexibility in adapting to changing market trends"
],
"planning_strategies": [
"Collaborate with city planners to integrate infrastructure with existing urban planning",
"Work closely with regulators to streamline certification processes and ensure compliance",
"Monitor market demand and adjust pricing strategy accordingly",
"Invest in staff training and technology to enhance customer experience",
"Stay agile and adaptable to respond to changing market dynamics"
],
"evaluation": {
"safety": 9,
"environmental impact": 9,
"reliability": 8,
"cost-effectiveness": 7
},
"reason": "Plan 2 excels in safety, environmental impact, and reliability with dedicated infrastructure, electric aircraft, and regulatory compliance. However, it may have higher upfront costs and limited flexibility compared to Plan 1."
}
}You can play around with your own example if you want to. I have shared the full code in my github repo — Tree of Thoughts with LLM using Langchain.
References:
- “Tree of Thoughts: Deliberate Problem Solving
with Large Language Models” — https://arxiv.org/pdf/2305.10601 - https://python.langchain.com/v0.2/docs/how_to/sequence/
- https://medium.com/@astropomeai/implementing-the-tree-of-thoughts-in-langchains-chain-f2ebc5864fac
Connect with me on LinkedIn — https://www.linkedin.com/in/meenakshisrinivasan
