Custom Tools in LangChain for Language Models

Introduction

Large language models and conversational agents have recently emerged as some of the most fascinating technologies. These models, equipped with a wide range of tools, offer limitless possibilities. The LangChain tools enable them to perform various tasks such as web searches, coding, mathematics, and more. In this blog, we will learn how to build custom tools in LangChain.

While pre-built tools can take us far, there are cases where customization or the creation of new tools becomes necessary. Custom Tools in LangChain are useful for fine-tuning the task as required, providing additional flexibility.

A Recap of Tools

In the context of large language models, tools are components that augment the model’s abilities. Think of them as entities that receive specific inputs and generate corresponding outputs. The model controls the input to a tool and expands its capabilities by receiving the output. Tools have various functions, but they all transform input data into meaningful output. They help large language models overcome limitations and tackle challenging tasks like mathematics.

Tools act as intermediaries, linking the large language model with specific tasks or domains. We can create and incorporate these tools into the model’s architecture, thereby enhancing its abilities in multiple ways. They enable the model to process and understand different types of data and generate appropriate responses.

One important aspect of tools is their ability to process specific inputs. For example, a mathematical tool can take in mathematical expressions or equations as input and compute the corresponding results. The language model alone lacks inherent mathematical capabilities. However, with a mathematical tool, it can perform calculations and provide accurate answers.

Tools can also extend the model’s capabilities in other domains, such as translation, summarization, sentiment analysis, or even generating code. We can create each custom tool in LangChain to process specific input data and generate relevant output, enabling the model to demonstrate expertise across multiple fields.

Custom Tools in LangChain

To illustrate the concept of tools, let’s consider a simple example of a circle circumference calculator tool. In this case, the large language model struggles with mathematical calculations, making it an ideal scenario for using a tool. The tool is defined using the LangChain tools library and inherits essential methods from a base tool class.

Let us see it in action:

from langchain.tools import BaseTool
from typing import Union
from math import pi

from typing import Any

# single input tool

class AreaofCircle(BaseTool):
    name = "Area of Circle"
    description = "Use this tool Calculate the area of a circle using radius"

    def _run(self, radius : Union[int,float]) -> Any:
        return pi * float(radius) ** 2 
    
    def _async_run(self) -> Any:
        raise NotImplementedError("This tool does not support async run")

Similar to how we define a tool custom tool containing the following parameters:

• Name: It denotes the name of the tool.
• Description: It helps LLM decide the purpose of the tool and it is passed in along with the prompt.

Now it contains two functions _run and _arun. _arun is used when you want to execute a function asynchronously. In contrast, _run is used to define the function that is executed synchronously.

in our case, it takes in a single parameter radius and union tells us that it can be either int or float. It will return the area of the circle. Now let us try and use it:

import os
from langchain.chat_models import ChatOpenAI
from langchain.chains.conversation.memory import ConversationBufferWindowMemory

OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') or 'YOUR API KEY'

# initialize LLM (we use ChatOpenAI because we'll later define a `chat` agent)
llm = ChatOpenAI(
    openai_api_key=OPENAI_API_KEY,
    temperature=0,
    model_name='gpt-3.5-turbo'
)
# initialize conversational memory
conversational_memory = ConversationBufferWindowMemory(
    memory_key='chat_history',
    k=5,
    return_messages=True
)

We initialize the llm instance and a chat conversational memory as we saw in our blogs it helps the model retain the last 5 conversations:

from langchain.agents import initialize_agent

tools = [AreaofCircle()]

agent = initialize_agent(
    agent='chat-conversational-react-description',
    tools=tools,
    llm=llm,
    verbose=True,
    max_iterations=3,
    early_stopping_method='generate',
    memory=conversational_memory
)

Next, we initialize a chat-conversational-react-description agent. Now let us see what this agent is chat refers to the fact that it can converse with the user while conversational stands and it is using conversational memory. Next comes functioning react tells us that it uses reasoning and action ideology to execute any task and to decide which tool to use it uses the description as denoted by the last word “description”.

Let us run this agent and see it in action:

agent("calculate the area of a circle with radius 5.9")

> Entering new AgentExecutor chain...
{
    "action": "Final Answer",
    "action_input": "The area of the circle is approximately 109.35 square units."
}

> Finished chain.
{'input': 'calculate the area of a circle with radius 5.9',
 'chat_history': [],
 'output': 'The area of the circle is approximately 109.35 square units.'}

It generates an answer of 109.35 which is correct but the approach it took is not since it directly jumped to answer generation as we can see from its action, “Final answer” Rather than using its inherent mathematical capabilities it should use tool for an accurate answer every time.

To achieve this we can try to see the prompt that it uses and see how it can be modified:

print(agent.agent.llm_chain.prompt.messages[0].prompt.template)

Assistant is a large language model trained by OpenAI.

Assistant is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, Assistant is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.

Assistant is constantly learning and improving, and its capabilities are constantly evolving. It is able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. Additionally, Assistant is able to generate its own text based on the input it receives, allowing it to engage in discussions and provide explanations and descriptions on a wide range of topics.

Overall, Assistant is a powerful system that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, Assistant is here to assist.

So now we can modify the existing prompt to instruct it to use tools whenever possible:

sys_msg = """Assistant is a large language model trained by OpenAI.

Assistant is designed to be able to assist with a wide range of tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. As a language model, Assistant is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.

Assistant is constantly learning and improving, and its capabilities are constantly evolving. It is able to process and understand large amounts of text, and can use this knowledge to provide accurate and informative responses to a wide range of questions. Additionally, Assistant is able to generate its own text based on the input it receives, allowing it to engage in discussions and provide explanations and descriptions on a wide range of topics.

Unfortunately, Assistant is terrible at maths. When provided with math questions, no matter how simple, assistant always refers to it's trusty tools and absolutely does NOT try to answer math questions by itself

Overall, Assistant is a powerful system that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics. Whether you need help with a specific question or just want to have a conversation about a particular topic, Assistant is here to assist.
"""

We added the line to instruct the model to use its tool as it is bad at doing maths, this will ensure that model uses tools rather than answering the questions itself.

new_prompt = agent.agent.create_prompt(
    system_message=sys_message,
    tools = tools,
)

agent.agent.llm_chain.prompt = new_prompt

agent("calculate the area of a circle with radius 5.9")

Now we set the agent prompt to the new prompt and execute the chain:

> Entering new AgentExecutor chain...
{
    "action": "Area of Circle",
    "action_input": "5.9"
}
Observation: 109.3588402714607
Thought:{
    "action": "Final Answer",
    "action_input": "The area of a circle with radius 5.9 is approximately 109.36 square units."
}

> Finished chain.
{'input': 'calculate the area of a circle with radius 5.9',
 'chat_history': [HumanMessage(content='calculate the area of a circle with radius 5.9', additional_kwargs={}, example=False),
  AIMessage(content='The area of the circle is approximately 109.35 square units.', additional_kwargs={}, example=False)],
 'output': 'The area of a circle with radius 5.9 is approximately 109.36 square units.'}

This time model chooses to use a tool to answer the query rather than using its inherent capabilities.

Custom Multi-Input Tool

In our previous example, we saw a tool that takes in a single input and generates a single output but what if we need a tool that can take in multiple inputs? So let us jump straight into it.

from typing import Optional

desc = (
    "use this tool when you need to calculate volume of a cylinder/sphere or cuboid. "
    "To use the tool you must pass base_radius and height for cylinder, radius_sphere for sphere and length, width and height for cuboid."
    "It takes the following parameters"
    "['radius_sphere ,'base_radius', 'height_cylinder', 'height_cuboid', 'length_cuboid', 'width_cuboid']"
)

class VolumeTool(BaseTool):

    name = "Volume calculator"
    description = desc

    def _run(
        self,
        radius_sphere: Optional[Union[int, float]] = None,
        base_radius: Optional[Union[int, float]] = None,
        height_cylinder: Optional[Union[int, float]] = None,
        height_cuboid: Optional[Union[int, float]] = None,
        length_cuboid: Optional[Union[int, float]] = None,
        width_cuboid: Optional[Union[int, float]] = None
    ):
        # check for the values we have been given
        if radius_sphere:
            return (4/3)*pi*(float(radius_sphere)**3)
        elif base_radius and height_cylinder:
            return pi*(float(base_radius)**2)*float(height_cylinder)
        elif height_cuboid and length_cuboid and width_cuboid:
            return float(height_cuboid)*float(length_cuboid)*float(width_cuboid)
        else:
            return "Could not calculate the volume. Need two or more of `radius_sphere`, `base_radius`, `height_cylinder`, `height_cuboid`, `length_cuboid`, `width_cuboid`."
        
    def _arun(self, query: str):
        raise NotImplementedError("This tool does not support async")

Let us start by understanding the description of this tool. This tool instructs LLM that it can use this tool to calculate volume, now we instruct LLM to pass in various parameters for different shapes base_radius and height for the cylinder, radius_sphere for the sphere and length, and width and height for the cuboid.

Following this, we instruct the model that functions takes the parameter which helps it in passing parameters now we need to initialize the tool and add it to the tool list along with the new prompt we created above:

tools = [VolumeTool()]

new_prompt = agent.agent.create_prompt(
    system_message=sys_msg,
    tools=tools
)

agent.agent.llm_chain.prompt = new_prompt
# update the agent tools
agent.tools = tools

Now let us execute the chain and see its output:

agent("calculate the volume of a sphere with radius 5.9")

> Entering new AgentExecutor chain...
{
    "action": "Volume calculator",
    "action_input": {
        "radius_sphere": 5.9
    }
}
Observation: 860.2895434688243
Thought:{
    "action": "Final Answer",
    "action_input": "The volume of a sphere with radius 5.9 is approximately 860.29 units cubed."
}

Conclusion

In this article, we learned how to build custom tools in LangChain with examples. Starting with a simple tool for calculating the circumference of a circle, we then moved on to a more complex tool that took in multiple parameters as well.