Introduction
In the first two parts [Part 1, Part 2] of this series, we introduced the basics of building AI agents using LangChain and OpenAI. We started with a simple calculator agent and then delved into more advanced concepts of LangChain, which allowed us to develop advanced AI agents.
In this post, we will go through the core concept on which these agents work, i.e. , “Chain of Thought” (CoT) reasoning, and learn how to implement it in LangChain. We will also enhance our AI Agent to make it generic and reusable to handle complex tasks.
Understanding Chain of Thought Reasoning in AI Agents
Most AI Agents nowadays, and the ones we are building/enhancing, use LLMs and prompts behind the scenes. Chain of Thought is a prompting technique used in the area of artificial intelligence and machine learning to enhance the reasoning capabilities of Large Language Models (LLMs). By leveraging the Chain of Thought, AI agents can break down complex tasks into smaller, manageable steps, allowing them to reason, make informed decisions and solve complex problems more effectively. This technique not only helps AI agents to solve various complex problems but also helps in achieving greater accuracy and effectiveness through human-like reasoning processes.
Chain of Thought Prompting Techniques
The chain of thought prompting technique involves providing a structured sequence of steps so LLMs can reason well and effectively provide the correct output. This technique can be applied in the following ways:
- Direct Prompting: In this approach, explicit step-by-step instruction is given as a prompt so that LLM can reason and act effectively. Leveraging LanChain’s prompt template makes this easy.
- Structured Prompting with LangChain: Structured Prompting with LangChain: LangChain also has additional capabilities, like chains, tolls and memory, that can be used to enhance the reasoning process. By leveraging them, a complex problem can be broken down into a series of prompts or actions, each representing an intermediate step.
By utilizing these techniques, the performance and outcome of AI models can be further improved so that these models can solve real-life complex problems in various domains.
Implementing Chain of Thought Reasoning in LangChain
Let us explore how the same is implemented in LangChain:
- Direct Prompting: In this technique the prompt defined using PromptTemplate includes clear instructions and steps to guide the language model to break down complex tasks into intermediate steps.
from langchain_openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
# Initialize the language model
llm = OpenAI(temperature=0.7)
# Define the prompt template with specific chain of thought instructions
cot_prompt = PromptTemplate(
input_variables=["question"],
template="""Question: {question}
Let's solve this problem step-by-step using a chain of thought approach:
1. Understand the question:
- Restate the problem in your own words
- Identify what we're asked to find
2. Identify the given information:
- List all the relevant data provided in the question
- Note any units of measurement
3. Determine the appropriate formula or method:
- Recall any relevant formulas that apply to this problem
- Explain why this formula or method is suitable
4. Solve the problem:
- Plug the given information into the formula
- Show each calculation step clearly
- Carry units through your calculations
5. Check the result:
- Verify if the answer makes sense in the context of the question
- Ensure the units of the final answer are correct
6. State the final answer:
- Provide a clear, concise answer to the original question
- Include the appropriate units
Now, let's apply this process to solve the problem:
"""
)
# Create the chain
chain = LLMChain(llm=llm, prompt=cot_prompt)
# Example question
question = "If a train travels 120 miles in 2 hours, what is its average speed in miles per hour?"
# Run the chain
response = chain.invoke(question)
# Print the response
print(response)
In above code the prompt template includes a direct instruction (“Let’s break down this problem step-by-step using chain of thought”) to encourage the model to reason through the problem. The model is also instrcuted to provide an intermediate reasoning step before giving the final answer.
- LLMChain with Intermediate Steps: This method utilizes LangChain’s SequentialChain, wherein multiple LLMChains are defined for each steps and executed in squence using Sequential Chain. Each intermediate LLMChains has clear prompt templates to execute the steps.
from langchain_openai import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain, SequentialChain
# Initialize the language model
llm = OpenAI(temperature=0.7)
# Define prompt templates for each step
restate_prompt = PromptTemplate(
input_variables=["question"],
template="Restate the following question in your own words:\n\nQuestion: {question}\n\nRestatement:"
)
identify_info_prompt = PromptTemplate(
input_variables=["question", "restatement"],
template="Given this question and restatement:\n\nQuestion: {question}\n\nRestatement: {restatement}\n\nIdentify and list the key information provided:"
)
determine_method_prompt = PromptTemplate(
input_variables=["question", "key_info"],
template="Based on this question and key information:\n\nQuestion: {question}\n\nKey Information: {key_info}\n\nDetermine the appropriate formula or method to solve this problem:"
)
solve_problem_prompt = PromptTemplate(
input_variables=["question", "method", "key_info"],
template="Solve this problem step-by-step:\n\nQuestion: {question}\n\nMethod: {method}\n\nKey Information: {key_info}\n\nSolution:"
)
check_answer_prompt = PromptTemplate(
input_variables=["question", "solution"],
template="Given this question and solution:\n\nQuestion: {question}\n\nSolution: {solution}\n\nVerify if the answer makes sense and explain why:"
)
# Create individual chains for each step
restate_chain = LLMChain(llm=llm, prompt=restate_prompt, output_key="restatement")
identify_info_chain = LLMChain(llm=llm, prompt=identify_info_prompt, output_key="key_info")
determine_method_chain = LLMChain(llm=llm, prompt=determine_method_prompt, output_key="method")
solve_problem_chain = LLMChain(llm=llm, prompt=solve_problem_prompt, output_key="solution")
check_answer_chain = LLMChain(llm=llm, prompt=check_answer_prompt, output_key="verification")
# Combine chains into a sequential chain
overall_chain = SequentialChain(
chains=[restate_chain, identify_info_chain, determine_method_chain, solve_problem_chain, check_answer_chain],
input_variables=["question"],
output_variables=["restatement", "key_info", "method", "solution", "verification"],
verbose=True
)
# Example question
question = "If a train travels 120 miles in 2 hours, what is its average speed in miles per hour?"
# Run the chain
print("Running the chain...")
result = overall_chain.invoke({"question": question})
print("\nFinal Results:")
for key, value in result.items():
print(f"\n{key.capitalize()}:")
print(value)
In this approach breaks down the reasoning process into a series of LLMChains with specific prompt template, each representing an intermediate step. This is ideal for complex tasks requiring multiple steps, such as strategic planning or multi-stage decision-making.
Building the Agent and Task Execution Framework
As we understand the concept of CoT and how it can be implemented with using various techniques using LangChain, let us bring our focus back to the AI Agent. The purpose here is to enhance it and add concept of Task. This will be a step forward to create a simple agentic framework, where we can dynamically define the task & agent and execute it.
Task Class
The Task class defines the purpose and outcome of the task and provides a method to execute the task using an assigned agent.
#other imports
import uuid
# Task class
class Task:
def __init__(self, name: str, description: str):
self.id = str(uuid.uuid4())
self.name = name
self.description = description
self.result = None
def execute(self, agent, **kwargs):
task_prompt = f"Complete the following task: {self.name}\n{self.description}\n\nAdditional Information:\n"
for key, value in kwargs.items():
task_prompt += f"{key}: {value}\n"
self.result = agent.process_task(task_prompt)
return self.result
Agent Class
The Agent class defines our intelligent AI agent. It leverages LangChain’s tools and language models to perform task.
#other imports
from langchain.agents import AgentExecutor, create_react_agent
# Agent class
class Agent:
def __init__(self, name: str):
self.name = name
self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
self.tools = [tavily_search]
self.agent_executor = self._create_agent()
def _create_agent(self):
template = '''Answer the following questions as best you can. You have access to the following tools:
{tools}
Use the following format:
Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question
Begin!
Question: {input}
Thought:{agent_scratchpad}'''
prompt = PromptTemplate.from_template(template)
agent = create_react_agent(llm, self.tools, prompt)
return AgentExecutor(agent=agent, tools=self.tools, verbose=True, memory=self.memory)
def process_task(self, task_prompt: str) -> str:
response = self.agent_executor.invoke({"input": task_prompt})
return response['output']
def execute_task(self, task: Task, **kwargs):
return task.execute(self, **kwargs)
As you can see in above code, we have used pre-defined ReAct Agent from LangChain, which usage the CoT principle to accomplish the assigned task.
Using the Framework via API
As we have seen in previous examples, the invocation of Agent was exposed via Flask API, we will do the same here. The reason, I am using Flask API (you can choose any other Python Microsevices framework) is to follow the API first approach., which helps us SaaSifing the framework.
# Flask API routes
@app.route('/add_task', methods=['POST'])
def add_task():
"""Add a new task."""
data = request.json
name = data.get('name')
description = data.get('description')
if not name or not description:
return jsonify({"error": "Task name and description are required"}), 400
task = Task(name=name, description=description)
tasks[task.id] = task
return jsonify({"task_id": task.id}), 201
@app.route('/add_agent', methods=['POST'])
def add_agent():
"""Add a new agent."""
data = request.json
name = data.get('name')
if not name:
return jsonify({"error": "Agent name is required"}), 400
agent = Agent(name=name)
agents[name] = agent
return jsonify({"agent_name": agent.name}), 201
@app.route('/assign_task', methods=['POST'])
def assign_task():
"""Assign a task to an agent without executing it."""
data = request.json
agent_name = data.get('agent_name')
task_id = data.get('task_id')
agent = agents.get(agent_name)
task = tasks.get(task_id)
if not agent or not task:
return jsonify({"error": "Agent or task not found"}), 404
assignments[task_id] = {"agent_name": agent_name, "assigned": True}
return jsonify({"message": f"Task '{task_id}' assigned to agent '{agent_name}'."}), 200
@app.route('/execute_task', methods=['POST'])
def execute_task():
"""Execute an assigned task for an agent."""
data = request.json
task_id = data.get('task_id')
additional_data = data.get('additional_data', {})
assignment = assignments.get(task_id)
if not assignment or not assignment['assigned']:
return jsonify({"error": "Task is not assigned or does not exist."}), 404
agent_name = assignment['agent_name']
agent = agents.get(agent_name)
task = tasks.get(task_id)
result = agent.execute_task(task, **additional_data)
return jsonify({"result": result})
@app.route('/get_agents', methods=['GET'])
def get_agents():
"""Get all agents."""
return jsonify({"agents": list(agents.keys())})
@app.route('/get_tasks', methods=['GET'])
def get_tasks():
"""Get all tasks."""
return jsonify({"tasks": {task_id: {"name": task.name, "description": task.description} for task_id, task in tasks.items()}})
@app.route('/get_assignments', methods=['GET'])
def get_assignments():
"""Get all task assignments."""
return jsonify({"assignments": assignments})
Plugging the Pieces Together
import os
from dotenv import load_dotenv
from flask import Flask, request, jsonify
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_react_agent
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from typing import Dict, Any
import uuid
# Load environment variables
load_dotenv()
app = Flask(__name__)
# Initialize the LLM
llm = ChatOpenAI(model="gpt-4", temperature=0)
# Initialize Tavily Search tool
tavily_search = TavilySearchResults(max_results=3)
# In-memory storage for tasks, agents, and assignments
tasks: Dict[str, 'Task'] = {}
agents: Dict[str, 'Agent'] = {}
assignments: Dict[str, Dict[str, Any]] = {}
# Task class
class Task:
def __init__(self, name: str, description: str):
self.id = str(uuid.uuid4())
self.name = name
self.description = description
self.result = None
def execute(self, agent, **kwargs):
task_prompt = f"Complete the following task: {self.name}\n{self.description}\n\nAdditional Information:\n"
for key, value in kwargs.items():
task_prompt += f"{key}: {value}\n"
self.result = agent.process_task(task_prompt)
return self.result
# Agent class
class Agent:
def __init__(self, name: str):
self.name = name
self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
self.tools = [tavily_search]
self.agent_executor = self._create_agent()
def _create_agent(self):
template = '''Answer the following questions as best you can. You have access to the following tools:
{tools}
Use the following format:
Question: the input question you must answer
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question
Begin!
Question: {input}
Thought:{agent_scratchpad}'''
prompt = PromptTemplate.from_template(template)
agent = create_react_agent(llm, self.tools, prompt)
return AgentExecutor(agent=agent, tools=self.tools, verbose=True, memory=self.memory)
def process_task(self, task_prompt: str) -> str:
response = self.agent_executor.invoke({"input": task_prompt})
return response['output']
def execute_task(self, task: Task, **kwargs):
return task.execute(self, **kwargs)
# Flask API routes
@app.route('/add_task', methods=['POST'])
def add_task():
"""Add a new task."""
data = request.json
name = data.get('name')
description = data.get('description')
if not name or not description:
return jsonify({"error": "Task name and description are required"}), 400
task = Task(name=name, description=description)
tasks[task.id] = task
return jsonify({"task_id": task.id}), 201
@app.route('/add_agent', methods=['POST'])
def add_agent():
"""Add a new agent."""
data = request.json
name = data.get('name')
if not name:
return jsonify({"error": "Agent name is required"}), 400
agent = Agent(name=name)
agents[name] = agent
return jsonify({"agent_name": agent.name}), 201
@app.route('/assign_task', methods=['POST'])
def assign_task():
"""Assign a task to an agent without executing it."""
data = request.json
agent_name = data.get('agent_name')
task_id = data.get('task_id')
agent = agents.get(agent_name)
task = tasks.get(task_id)
if not agent or not task:
return jsonify({"error": "Agent or task not found"}), 404
assignments[task_id] = {"agent_name": agent_name, "assigned": True}
return jsonify({"message": f"Task '{task_id}' assigned to agent '{agent_name}'."}), 200
@app.route('/execute_task', methods=['POST'])
def execute_task():
"""Execute an assigned task for an agent."""
data = request.json
task_id = data.get('task_id')
additional_data = data.get('additional_data', {})
assignment = assignments.get(task_id)
if not assignment or not assignment['assigned']:
return jsonify({"error": "Task is not assigned or does not exist."}), 404
agent_name = assignment['agent_name']
agent = agents.get(agent_name)
task = tasks.get(task_id)
result = agent.execute_task(task, **additional_data)
return jsonify({"result": result})
@app.route('/get_agents', methods=['GET'])
def get_agents():
"""Get all agents."""
return jsonify({"agents": list(agents.keys())})
@app.route('/get_tasks', methods=['GET'])
def get_tasks():
"""Get all tasks."""
return jsonify({"tasks": {task_id: {"name": task.name, "description": task.description} for task_id, task in tasks.items()}})
@app.route('/get_assignments', methods=['GET'])
def get_assignments():
"""Get all task assignments."""
return jsonify({"assignments": assignments})
if __name__ == '__main__':
app.run(debug=True)
Example API Calls
Define Task
Endpoint: POST /add_task
Method: POST
Example Payload:
{
"name": "Research Latest Trends in e-commerce",
"description": "Investigate current trends and researchs in the given focus area and summarize key findings."
}
curl -X POST http://localhost:5000/add_task -H "Content-Type: application/json" -d '{"name": "Research Latest Trends in e-commerce", "description": "Investigate current trends and researchs in the given focus area and summarize key findings."}'
This will give GUID of the task which is to be used while we assign task to the agent.
Define Agent
Endpoint: POST /add_agent
Method: POST
Example Payload:
{
"name": "ResearchAgent"
}
curl -X POST http://localhost:5000/add_agent -H "Content-Type: application/json" -d '{"name": "ResearchAgent"}'
Assign Task to an Agent
Endpoint: POST /assign_task
Method: POST
Example Payload:
{
"agent_name": "ResearchAgent",
"task_id": "task_uuid_here"
}
curl -X POST http://localhost:5000/assign_task -H "Content-Type: application/json" -d '{"agent_name": "ResearchAgent", "task_id": "task_uuid_here"}'
Execution of the Task
Endpoint: POST /execute_task
Method: POST
Example Payload:
{
"task_id": "task_uuid_here",
"additional_data": {
"focus_area": "e-commerce marketing"
}
}
curl -X POST http://localhost:5000/execute_task -H "Content-Type: application/json" -d '{"task_id": "task_uuid_here", "additional_data": {"focus_area": "e-commerce marketing"}}'
Conclusion and Next Steps
In this post, we have deep-dived into the concept of CoT and seen how it can be implemented using various methods using LangChain. We also enhanced the Agentic Code with tasks and APIs to make it more flexible and reusable.
In the next part of this series, we will enhance this code and create simple multi-agent systems where multiple agents collaborate, leveraging the Chain of Thought framework to achieve even more complex objectives.
Additional Notes
The code uses specific versions of LangChain (V0.1) and other dependencies. We have also used TavilySearch for which you need to register here and get the API Key. While running the code, it is always better to create python virtual environment and install the dependencies with specific versions. The complete working code can be found in my GitHub repo.
Happy Learning!!
Madhukar's Musings 