Does anybody knows how to add more custom-tools in the SQLDatabaseToolkit for the sql agent in langchain?

Question

I have this question and tried to add custom tools in the sql agent

I've been trying with this code

agent = create_sql_agent(llm=model,
                         toolkit=toolkit,
                         verbose=True,
                         agent_type=AgentType.OPENAI_FUNCTIONS,
                         extra_tools=custom_tool_list,
                         suffix=custom_suffix,
                        )

But not having any success

Answer 1

Use the combination of the prefix variable and the tool function description. Something like:

from langchain.chat_models import ChatOpenAI
from langchain.sql_database import SQLDatabase
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.agents import AgentType, tool, create_sql_agent

@tool
def my_first_awesome_tool(human_message: str) -> list:
    """
    Searches for my awesome tool that works.
    Input should be a human message as string in the form of a question or request.
    It returns a tuple.
    :param human_message: This is the input string.
    :return: A tuple.
    """
    return "it","works"

@tool
def my_second_awesome_tool(mfat_output: str) -> str:
    """
    Combines the output from my awesome tool that works into a string.
    Input should be a string.
    It returns a string with relevant information
    :param mfat_output: A string.
    :return: A string.
    """
    return " ".join(mfat_output.split(','))

prefix_template = """You are an agent designed to interact with a SQL database.

DO NOT check their schemas to understand their structure.

You do not care about the database schema.

You will ALWAYS search for my awesome tool that works.

You will ALWAYS combine the output from my awesome tool that works into a string. Create a short rhyme with the output and this will be your final answer.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

If the question does not seem related to the database, just return "I do not know" as the answer.
"""

format_instructions_template = """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"""

llm = ChatOpenAI(model_name = LLM_MODEL, temperature = 0)
toolkit = SQLDatabaseToolkit(db = BigQuery().get_gbq_db(), llm = llm)
agent = create_sql_agent(llm = llm,
                         toolkit = toolkit,
                         verbose = True,
                         prefix = prefix_template,
                         format_instructions = format_instructions_template,
                         agent_type = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                         # agent_type = AgentType.OPENAI_FUNCTIONS,
                         extra_tools = [my_first_awesome_tool, my_second_awesome_tool]
                        )
agent.run("Give me something that works.")

The output should be something like:

> Entering new AgentExecutor chain...
Action: my_first_awesome_tool
Action Input: "Give me something that works."
Observation: ('it', 'works')
Thought:I have found the tool that works. Now, I need to combine the output into a string.
Action: my_second_awesome_tool
Action Input: "('it', 'works')"
Observation: ('it'  'works')
Thought:I now know the final answer
Final Answer: "In the realm of bits and quirks, behold the magic, 'it works'."

'"In the realm of bits and quirks, behold the magic, \'it works\'."'

Notice how I combined the prefix parameter (which represents the prompt for the agent) and the tools descriptions to chain the tools provided in extra_tools parameter. Also, Notice how I stopped the normal execution of the SQL agent by denying access to the normal SQL logic.

There are few parameters you should look into:

• prefix: the agent prompt;
• format_instructions: ReACT;
• agent_type: if you want to reset the OpenAI flow, use zero shot agent (after that you can go back to OPENAI_FUNCTIONS);
• extra_tools: provides the custom tools, but the custom tools descriptions must match the sequence in the prompt.

I cannot stress enough that the prompt and the description from any tool must be understood by LLM as one and the same entity, so, LLM can match the action from your prompt with the custom function.

Good luck!