[LangChain] Custom Tool

This page is a note about how to custom tool with LangChain

Python=3.10.0

langchain=0.0.320

Methods

tool: Use decorator @tool

Tool: Use Tool.from_function

StructuredTool: Use StructuredTool.from_function

BaseTool: Inherit the BaseTool class

Examples

1. Import Modules

from langchain.llms import AzureOpenAI

from langchain.tools import tool, Tool, BaseTool, StructuredTool
from langchain.tools.base import ToolException
from langchain.agents import AgentType, initialize_agent

2. Custom Tool

# Use decorator "@tool" to get Tool object
@tool
def add(a: int, b: int) -> int:
    """Add two inputs, and return it"""
    return a + b

def subtract(query: str) -> int:
    a, b = query.split(",")
    return int(a) - int(b)

def multiply(a: int, b: int) -> int:
    return a * b

# Inherit "BaseTool" to customize the Your_Tool class
class Divide(BaseTool):
    name = "divide"
    description = """divide two inputs, and return it"""

    def _run(self, a: int, b: int) -> int:
        """Use the tool"""
        return a / b

    def _arun(self, a: int, b: int) -> int:
        """Use the tool asynchronously."""
        raise NotImplementedError("Does not support async.")

tools = [
    add,
    # Use "Tool.from_function" to get Tool object
    Tool.from_function(
        func = subtract,
        name = "subtract",
        description = """subtract two numbers and return it.
            The input should be separated by ",".
            Example: 5,7
        """
    ),
    # Use "StructuredTool.from_function" to get Tool object
    StructuredTool.from_function(
        func = multiply,
        name = "multiply",
        description = """Multiply two inputs, and return it"""
    ),
    # Create Tool instance
    Divide()
    ]

3. Initialize agent and run

AgentType.ZERO_SHOT_REACT_DESCRIPTION 不支援Multi-Input 的Function，

因此建議使用支援Multi-input 的 AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION。

# LLM
llm = AzureOpenAI(
    temperature=0,
    engine = "OpenAImodel",
    model = "gpt-3.5-turbo",
    verbose=True,
)

agent = initialize_agent(
    tools, 
    llm, 
    # agent = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    agent = AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose = True,
    handle_parsing_errors = True,
    return_direct = True
)

result = agent.run("Calculate 1 + 2 - 3 * 4 / 6, step by step")
print(f"Result: {result}")