skip to Main Content

Comet is now available natively within AWS SageMaker!

Learn More

Enhancing LangChain Agents with Custom Tools

How to Create and Implement Custom Tools in LangChain

Enhancing LangChain Agents with Custom Tools, Comet + CometLLM
Photo by Todd Quackenbush on Unsplash

Preliminaries

%%capture
!pip install langchain openai duckduckgo-search youtube_search wikipedia

import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter Your OpenAI API Key:")

Tool Input Schema

This is the pattern for loading tools in LangChain:

from langchain.agents import load_tools

tool_names = ['your', 'list', 'of', 'tools']
tools = load_tools(tool_names)

Some tools might require a base LLM to initialize them. You would change it to the following:

tools = load_tools(tool_names, llm=llm)

Defining Custom Tools

Custom tools in LangChain are defined by the user to perform specific tasks or operations not provided by the native tools in the LangChain toolkit.

They allow users to extend the functionality of LangChain and tailor it to their specific needs. There are several reasons why you might need custom tools in LangChain:

  • Specific Use Cases: Custom tools allow you to address specific use cases or requirements not covered by the native tools. You can define tools that perform complex calculations, interact with external APIs, manipulate data uniquely, or integrate with other systems.
  • Integration: Custom tools can integrate LangChain with other systems or tools. You can define tools that interact with external APIs, databases, or other software components to exchange data or perform actions.
  • Automation: Custom tools enable you to automate repetitive tasks or workflows. By defining a custom tool, you can encapsulate a series of steps or operations into a single tool that can be easily reused and executed with different inputs.
  • Domain-Specific Functionality: Custom tools allow you to incorporate domain-specific functionality into LangChain. For example, suppose you are working in the finance industry. In that case, you can define custom financial calculations or data analysis tools tailored to your needs.

Want to learn how to build modern software with LLMs using the newest tools and techniques in the field? Check out this free LLMOps course from industry expert Elvis Saravia of DAIR.AI!


Creating a custom tool in LangChain

To define a custom tool in LangChain, you can use the Tool.from_function() method or subclass the BaseTool class.

The Tool.from_function() method lets you quickly create a tool from a simple function. Subclassing the BaseTool class provides more control over the tool’s behaviour and defines custom instance variables or propagates callbacks. Once defined, custom tools can be added to the LangChain agent using the initialize_agent() method.

The agent can then execute the custom and native tools in a chain to perform complex workflows or tasks. Tool consists of several components:

  • name (str), is required and must be unique within a set of tools provided to an agent
  • description (str), is optional but recommended, as an agent uses it to determine tool use
  • return_direct (bool), defaults to False
  • args_schema (Pydantic BaseModel), is optional but recommended and can be used to provide more information (e.g., few-shot examples) or validation for expected parameters.

Before we create a tool, let’s set up our environment:

from langchain import LLMMathChain
from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.tools import BaseTool, StructuredTool, Tool, tool, DuckDuckGoSearchRun

from pydantic import BaseModel, Field

llm = ChatOpenAI(temperature=0)

A simple tool

The simplest tool is basically strings in, strings out. You pass in a query (text string), and you get a result (text string).

There are two ways you can do this:

1) Use the Tool dataclass

2) Subclass the BaseTool class

Using the Tool dataclass

This wraps functions that accept a single string input and return a string output.

search = DuckDuckGoSearchRun()

search_tool = Tool.from_function(
    func=search.run,
    name="Search",
    description="useful for when you need to search the internet for information"
)

llm_math_chain = LLMMathChain(llm=llm, verbose=True)

math_tool = Tool.from_function(
    func=llm_math_chain.run,
    name="Calculator",
    description="Useful for when you are asked to perform math calculations"
)

tools = [search_tool, math_tool]

agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

agent.run("Search for the Sacramento Kings win-loss record for the entire 2022-2023 season. Add games won and games lost and multiply that result by 2")
> Entering new AgentExecutor chain...
I need to search for the Sacramento Kings win-loss record for the 2022-2023 season. Then I need to add the games won and games lost and multiply that result by 2.
Action: Search
Action Input: "Sacramento Kings win-loss record 2022-2023 season"
Observation: Sacramento Kings 2022-2023 Schedule and Results All the scores of the games played by the Kings in the NBA 2022-23 Regular Season and the 2023 Playoffs, with teams, results, winning and losing streaks, partial and final records and evolution of the playoff series. Land Of Basketball.com Little did we know we'd be in for a historic night. With a 176-175 double-overtime win, the Kings set a new franchise single-game scoring record. The game also put both teams in the history books as the second-highest-scoring game in NBA history. 5️⃣ Monk and Fox Take Over LA Schedule & Results Transactions More 2022-23 Kings Pages Game Log Splits Lineups On/Off Starting Lineups Depth Charts Referees On this page: Roster Assistant Coaches and Staff Team and Opponent Stats Team Misc Per Game Totals Per 36 Minutes Per 100 Poss Advanced Adjusted Shooting Sacramento Kings record 2022-23. Explore Sacramento Kings all-time NBA seasons record Splits including Home Games, Road Games, Last 10 Games and much more. Expected W-L: 47-35 (8th of 30) Preseason Odds: Championship +50000, Over-Under 34.5. Arena: Golden 1 Center Attendance: 715,491 (20th of 30) NBA 2023 Playoffs: Lost NBA Western Conference First Round (3-4) versus Golden State Warriors ( Series Stats) Pacific Div Champs. Kings Franchise Index. Roster & Stats.
Thought:I found the Sacramento Kings' record for the 2022-2023 season, but I need to add the games won and games lost and multiply that result by 2.
Action: Search
Action Input: "Sacramento Kings record 2022-2023 season"
Observation: Sacramento Kings 2022-2023 Schedule and Results All the scores of the games played by the Kings in the NBA 2022-23 Regular Season and the 2023 Playoffs, with teams, results, winning and losing streaks, partial and final records and evolution of the playoff series. Land Of Basketball.com 2022-23 season schedule Regular season 2022-23 game log Total: 48-34 (home: 23-18; road: 25-16) October: 2-4 (home: 1-3; road: 1-1) November : 9-5 (home: 6-1; road: 3-4) December: 8-6 (home: 4-3; road: 4-3) January : 9-6 (home: 5-4; road: 4-2) February : 8-4 (home: 2-1; road: 6-3) March : 11-5 (home: 5-4; road: 6-1) Schedule & Results Transactions More 2022-23 Kings Pages Game Log Splits Lineups On/Off Starting Lineups Depth Charts Referees An ad blocker has likely prevented this video content from loading. Please disable your ad blocker to view the video content. Regular Season Schedule Playoffs Schedule Regular Season Share & Export Roster & Stats Schedule & Results Transactions More 2022-23 Kings Pages Game Log Splits Lineups On/Off Starting Lineups Depth Charts Referees On this page: Roster Assistant Coaches and Staff Team and Opponent Stats Team Misc Per Game Totals Per 36 Minutes Per 100 Poss Advanced Adjusted Shooting And when the final buzzer sounded on the regular season, with 9,846 points, the team set a new single-season franchise record for points scored, surpassing the 1961-62 Cincinnati Royals by one point. Sacramento finishes Q1 at DEN with 43 points and has now scored 9,846 points this season.
Thought:I have found the Sacramento Kings' win-loss record for the 2022-2023 season. Now I need to add the games won and games lost and multiply that result by 2.
Action: Calculator
Action Input: 48 (games won) + 34 (games lost)

> Entering new LLMMathChain chain...
48 (games won) + 34 (games lost)```text
48 + 34
```
...numexpr.evaluate("48 + 34")...

Answer: 82
> Finished chain.

Observation: Answer: 82
Thought:I now know the final answer. 
Final Answer: 82

> Finished chain.
82

Subclassing from the BaseTool class

When creating a custom tool in LangChain, you may subclass from the BaseTool class when you need more control over the tool’s behaviour and want to define custom instance variables or propagate callbacks.

Subclassing from the BaseTool class allows you to extend the tool’s functionality and customize its behaviour according to your requirements.

By subclassing, you can define additional methods, override existing methods, and add custom instance variables specific to your tool.

Subclassing from the BaseTool class enables you to propagate callback functions executed at specific points during the tool’s execution.

Callbacks can perform additional actions or modify the tool’s behaviour based on certain conditions.

See the documentation for a code example.

Using the tool decorator

The @tool decorator in LangChain creates a Custom Tool from a simple function.

Use the @tool decorator to quickly create a Tool without needing advanced customization or defining custom instance variables. The @tool decorator differs from the Tool dataclass in providing a more lightweight and function-based approach to creating a Tool. It allows you to define the Tool’s functionality directly within the decorated function, making it easier to encapsulate specific tasks or operations.

To determine whether to use the @tool decorator or the Tool data class, consider the level of customization and control you require for your Tool.

If you need advanced customization and want to define custom instance variables or propagate callbacks, then using the Tool dataclass and subclassing from the BaseTool class would be more suitable. However, using the decorator would be a good choice if you prefer a more straightforward and concise approach without extensive customization.

The @tool decorator provides a quick and easy way to create a Custom Tool from a function. At the same time, the Tool dataclass offers more flexibility and control over the Tool’s behavior. The choice between the decorator and the Tool dataclass depends on the customization you need for your Tool.

from langchain.tools import tool

@tool("lower_case", return_direct=True)
def to_lower_case(input:str) -> str:
  """Returns the input as all lower case."""
  return input.lower()
StructuredTool(name='lower_case', description='lower_case(input: str) -> str - Returns the input as all lower case.', args_schema=<class 'pydantic.v1.main.lower_caseSchemaSchema'>, return_direct=True, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, func=<function to_lower_case at 0x7dd2a05f2710>, coroutine=None)

Let’s add this tool to our existing list, which was defined above:

tools.append(to_lower_case)

agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True
)

agent.run("What is the Apple TV show Foundation about? Return the answer in all lower case")
> Entering new AgentExecutor chain...
I need to search for information about the Apple TV show Foundation and then return the answer in all lower case.
Action: Search
Action Input: "Apple TV show Foundation"
Observation: Foundation season 2 will debut on Apple TV Plus on Friday, July 14. Its release date was revealed alongside the series' first official trailer, and came four months after Apple confirmed a... Foundation Episodes & Images Cast & Crew Trailers & Videos Watch on Apple TV+ Episodes & Images S2, E8 • September 1, 2023 The Last Empress Enjoiner Rue confides in Dusk about her distrust of Demerzel. Hober Mallow pulls a daring move. Day sets course for Terminus and the Foundation. published 5 September 2023 This is a high-stakes dogfight we can get behind Get ready for an explosive ninth episode of Foundation season 2 (Image credit: Apple TV Plus) Foundation season 2... published 13 July 2023 Exclusive: 'I don't know who started that rumour' 'Brother Day, there are rumors we're getting a third season. Is it true?' (Image credit: Apple TV Plus) Has Apple... Features Foundation Recap: What Happened in Season 1? Need a refresher regarding what happened in Foundation season 1? We've got a year-by-year recap of events before you jump into season 2! By...
Thought:I need to extract the relevant information about the Apple TV show Foundation from the search results.
Action: Search
Action Input: "Apple TV show Foundation"
Observation: Foundation season 2 will debut on Apple TV Plus on Friday, July 14. Its release date was revealed alongside the series' first official trailer, and came four months after Apple confirmed a... Foundation Episodes & Images Cast & Crew Trailers & Videos Watch on Apple TV+ Episodes & Images S2, E8 • September 1, 2023 The Last Empress Enjoiner Rue confides in Dusk about her distrust of Demerzel. Hober Mallow pulls a daring move. Day sets course for Terminus and the Foundation. published 5 September 2023 This is a high-stakes dogfight we can get behind Get ready for an explosive ninth episode of Foundation season 2 (Image credit: Apple TV Plus) Foundation season 2... "Foundation: The Official Podcast," a companion podcast to the drama series, is hosted by Jason Concepcion and Goyer, and produced by Pineapple Street Studios with Max Linsky, Jenna Weiss-Berman and Bari Finkel serving as executive producers. The entire first season is available now on Apple Podcasts. This is the ambitious scope of Foundation, the Apple TV Plus show adapted from the Isaac Asimov series of the same name covering centuries of in-universe history, spanning years and...
Thought:I have extracted the relevant information about the Apple TV show Foundation. Now I need to return the answer in all lower case.
Action: lower_case
Action Input: "Foundation season 2 will debut on Apple TV Plus on Friday, July 14."
Observation: foundation season 2 will debut on apple tv plus on friday, july 14.


> Finished chain.
foundation season 2 will debut on apple tv plus on friday, july 14.

Modifying existing tools

You can find the available native tools in LangChain here and look at the dictionary _EXTRA_OPTIONAL_TOOLS for the key of the tool.

from langchain.agents import load_tools
tools = load_tools(["wikipedia"], llm=llm)

tools[0]
WikipediaQueryRun(name='Wikipedia', description='A wrapper around Wikipedia. Useful for when you need to answer general questions about people, places, companies, facts, historical events, or other subjects. Input should be a search query.', args_schema=None, return_direct=False, verbose=False, callbacks=None, callback_manager=None, tags=None, metadata=None, handle_tool_error=False, api_wrapper=WikipediaAPIWrapper(wiki_client=<module 'wikipedia' from '/usr/local/lib/python3.10/dist-packages/wikipedia/__init__.py'>, top_k_results=3, lang='en', load_all_available_meta=False, doc_content_chars_max=4000))

And you can modify the fields above. For example, let’s change the description.

tools[0].description
A wrapper around Wikipedia. Useful for when you need to answer general questions about people, places, companies, facts, historical events, or other subjects. Input should be a search query.
tools[0].description = "Use this when you are asked to search for people, places, or things."

Prioritizing tools over other tools

This is pretty easy. All you need to do is add a description like:

Use this tool more than regular search if you are asked factual information about a person, place, or thing

or

Use this tool more than the Wikipedia tool if you are asked about current events, recent information, or news

Handling Errors

The agent will stop executing when a tool encounters an error, and the exception is missed.

If you want the agent to continue execution, you can raise a ToolException and set handle_tool_error accordingly. When ToolException is thrown, the agent will not stop working but will handle the exception according to the handle_tool_error variable of the tool. The processing result will be returned to the agent as observation and printed in red.

You can set handle_tool_error to True, set it as a unified string value, or set it as a function.

If set as a function, it should take a ToolException as a parameter and return a str value. Please note that only raising a ToolException won’t be effective. You need to first set the handle_tool_error of the tool because its default value is False.

from langchain.tools.base import ToolException

from langchain.tools import DuckDuckGoSearchRun
from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.tools import Tool

from langchain.chat_models import ChatOpenAI


def _handle_error(error: ToolException) -> str:
    return (
        "The following errors occurred during tool execution:"
        + error.args[0]
        + "Please try another tool."
    )


def search_tool1(s: str):
    raise ToolException("The search tool1 is not available.")


def search_tool2(s: str):
    raise ToolException("The search tool2 is not available.")


search_tool3 = DuckDuckGoSearchRun()

description = "useful for when you need to answer questions about current events.You should give priority to using it."

tools = [
    Tool.from_function(
        func=search_tool1,
        name="Search_tool1",
        description=description,
        handle_tool_error=True,
    ),
    Tool.from_function(
        func=search_tool2,
        name="Search_tool2",
        description=description,
        handle_tool_error=_handle_error,
    ),
    Tool.from_function(
        func=search_tool3.run,
        name="Search_tool3",
        description="useful for when you need to answer questions about current events",
    ),
]

agent = initialize_agent(
    tools,
    ChatOpenAI(temperature=0),
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
)

agent.run("What is the Capital of Manitoba?")
> Entering new AgentExecutor chain...
I should use a search tool to find the answer to this question.
Action: Search_tool1
Action Input: "Capital of Manitoba"
Observation: The search tool1 is not available.
Thought:I should try using another search tool to find the answer.
Action: Search_tool2
Action Input: "Capital of Manitoba"
Observation: The following errors occurred during tool execution:The search tool2 is not available.Please try another tool.
Thought:I should try using the last search tool to find the answer.
Action: Search_tool3
Action Input: "Capital of Manitoba"
Observation: Capital: Winnipeg Population: (2021) 1,342,153 Date Of Admission: 1870 Provincial Motto: "Glorious and Free" Manitoba's capital and largest city is Winnipeg, the seventh most populous municipality in Canada. Winnipeg is the seat of government, home to the Legislative Assembly of Manitoba and the Provincial Court. Winnipeg, city, capital (1870) of Manitoba, Canada. It lies at the confluence of the Red and Assiniboine rivers, 40 miles (65 km) southwest of Lake Winnipeg and 60 miles (95 km) north of the U.S. state of Minnesota. Winnipeg is the economic and cultural centre of Manitoba and is at the heart of the most populous metropolitan area in central Canada. The capital of Manitoba is Winnipeg. Geographic and Economic Importance . Winnipeg is the capital of Manitoba which is the fifth most populous province of Canada. Winnipeg is situated near the longitudinal center of North America. Since it is in the middle of Canada's land mass, all rail and road traffic between Canada's East and West ... Key Facts Manitoba, a central Canadian province, boasts a diverse and fascinating geography that warrants further exploration. This comprehensive analysis delves into the province's major geographical features, including its location, climate, landforms, bodies of water, and natural resources.
Thought:I now know the final answer.
Final Answer: The capital of Manitoba is Winnipeg.

> Finished chain.
The capital of Manitoba is Winnipeg.

Conclusion

In closing, the post has navigated through the sophisticated features of LangChain, showcasing how custom tools and memory elevate an agent’s capabilities.

From expanding functionality with specific use cases to seamless integration and automation, LangChain agents are reshaping how we interact with AI systems. With a focus on practical implementation, the blog provides a step-by-step guide to leveraging these tools, ensuring that your AI agents are more efficient and intelligent in processing and decision-making tasks.

Whether for personal use or scaling enterprise solutions, LangChain’s flexible framework offers a powerful ally in the evolving landscape of AI technology.


Harpreet Sahota

Back To Top