Langchain is an awesome library that allows you to take LLM usage to the next level.
The coolest part is the ability to give control over some tools to the LLM. In langchain this idea comes in the form of agents.
However, even though there are lots of examples on how to use agents there is no clear way of how to combine them and reroute certain questions to different agents.
In this short tutorial I will present you how you can do that. We will create a router chain that can redirect questions between three agents.
One will be a CSV agent that answers questions about Spotify Top Songs from 2023, another is capable of using custom made tools to answer our questions, and the last one can answer simple math questions.
You can download the dataset used in the code below here.
In order to run the code examples below you will need to setup a python environment with langchain installed.
At the time of writing there is a bug with the latest version of langchain, so you may need to pin it to version `0.0.301`.
First we will need to make sure we import all the necessary modules.
import datetime
from langchain.agents import create_csv_agent
from langchain.agents import AgentType, initialize_agent
from langchain.chains import LLMMathChain
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import RunnableMap
from langchain.tools import Tool
Next we will want to create our agents.
Make sure you create an instance of the LLM that you want to use within the agents and the router chain. I recommend using OpenAI GPT-4 either through OpenAI's API or through your own deployment in Azure.
GPT-3.5 may still work, but generally the agents behave much better with GPT-4.
llm = ChatOpenAI(
model_name="gpt-4-0613" # Choose any OpenAI model you like
openai_api_key="YOUR_KEY", # You can set this through ENV var too - just set OPENAI_API_KEY
temperature=0., # Make sure your code tries to do the same thing every time
) # Alternatively, you can use AzureChatOpenAI with your custom Azure deployment
# I found that in order to get the csv agent to return what we want we need to modify the prefix to force it to evaluate the result variable.
csv_agent = create_csv_agent(
llm, # Same model will be used in many places
"spotify-2023.csv", # This is the name of the csv that you will download from kaggle
verbose=True,
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
prefix="""
You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
If you are using python_repl_ast and you know the result to be variable make sure you add a new python_repl_ast command at the end with a statement in which you print the result variable to the "Action Input:" part.
You should use the tools below to answer the question posed of you:""",
pandas_kwargs={
"encoding": "latin-1",
"sep": ",",
} # Necessary kwargs to parse the kaggle csv correctly on windows.
)
time_tool = Tool.from_function(
func=lambda input: datetime.datetime.now(),
name="Time Now",
description="useful when you need to know what time it is right now, so that you can process words like yesterday",
)
custom_tool = Tool.from_function(
func=lambda input: "42",
name="Answer to Hardest question",
description="useful for returning the answer to hardest question",
)
tools = [
time_tool,
custom_tool,
]
custom_agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
llm_math_chain = LLMMathChain(llm=llm)
math_tool = Tool.from_function(
func=llm_math_chain.run,
name="Math Calculator",
description="useful for when you need to do simple math calculations on real numbers",
return_direct=True,
)
math_agent = initialize_agent([math_tool], llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
Okay, so now that we have created our agents we will want to create a router chain in front of them, so that queries about spotify can go to the csv agent and other queries can go to the custom agent.
To do that we will need to create a simple chain with a custom parser/select function.
prompt = PromptTemplate.from_template("""If the question is about spotify, respond with `SPOTIFY`. If it is about math, respond with `MATH`. Otherwise, respond `OTHER`.
Question: {question}""")
router_chain = prompt | llm | StrOutputParser()
def select_chain(output):
if output["action"] == "SPOTIFY":
return csv_agent
elif output["action"] == "MATH":
return math_agent
elif output["action"] == "OTHER":
return custom_agent
else:
raise ValueError
chain = RunnableMap({
"action": router_chain,
"input": lambda x: x["question"]
}) | select_chain
Alright, now that we have created our router chain let us run some examples.
chain.invoke({"question": "What is the answer to the hardest question?"})
This returns as expected:
> Entering new AgentExecutor chain...
The question is asking for the answer to the hardest question. I have a tool that can provide this.
Action: Answer to Hardest question
Action Input: None
Observation: 42
Thought:I now know the final answer.
Final Answer: 42
> Finished chain.
What about some other ones?
chain.invoke({"question": "What date was two weeks ago?"})
and we get this as a result:
> Entering new AgentExecutor chain...
I need to know the current date to calculate the date two weeks ago.
Action: Time Now
Action Input: Current date
Observation: 2023-10-10 00:30:23.223340
Thought:Now that I know today's date, I can calculate the date two weeks ago.
Final Answer: 2023-09-26
> Finished chain.
Now to the math agent:
chain.invoke({"question": "What is 2**7?"})
Again, correct answer:
> Entering new AgentExecutor chain...
This is a simple math calculation. I need to calculate the power of 2 to the 7th degree.
Action: Math Calculator
Action Input: 2**7
Observation: Answer: 128
> Finished chain.
Okay, but what about our CSV agent. Can we route our questions to it too?
chain.invoke({"question": "Which song was the 3rd most famous song on spotify that was released in 2023?"})
And we get a correct response!
> Entering new AgentExecutor chain...
Thought: To find the 3rd most famous song on Spotify that was released in 2023, I need to filter the dataframe to only include songs released in 2023, then sort by the 'streams' column in descending order, and finally select the third row.
Action: python_repl_ast
Action Input: df_2023 = df[df['released_year'] == 2023]
df_2023_sorted = df_2023.sort_values(by='streams', ascending=False)
third_most_famous_song = df_2023_sorted.iloc[2]['track_name']
third_most_famous_song
Observation: People Pt.2 (feat. IU)
Thought:I now know the final answer
Final Answer: The 3rd most famous song on Spotify that was released in 2023 is "People Pt.2 (feat. IU)".
> Finished chain.
As you can see in the examples above GPT-4 with langchain can be quite powerful. The key however is in prompt engineering.
Initially, I could not get langchain CSV agent to work out of the box and ended up specifying a new line in the prompt to self correct for the lack of result value.
Fortunately, langchain gives you a lot of flexibility, so you can make these changes pretty easily. It is definitely not a finished product, but I consider it to be a great set of examples that you have to tune yourself to get the result you want.