Abstract
LangChain is a powerful framework that includes a variety of tools, including the agent_toolkits. In this article, we'll use the SQLDatabaseToolkit to interact with SingleStoreDB by making a request in English.
The notebook file used in this article is available on GitHub.
Introduction
A recent blog post described an internal SingleStore Demothon. One of the cool demos showed how to use LangChain's SQLDatabaseToolkit with SingleStoreDB. In this short article, we'll replicate the demo.
Create a SingleStoreDB Cloud account
A previous article showed the steps to create a free SingleStoreDB Cloud account. We'll use LangChain Demo Group as our Workspace Group Name and langchain-demo as our Workspace Name.
Create the TPC-H database
Once the workspace is available, we'll select the Tutorials (identified by the white question mark on a black background on the bottom right-hand side of the web page) and then choose Tutorials > Try with Sample Data > Benchmark Dataset - TPC-H. We'll check that the langchain-demo workspace is selected and then click the button to ▶ Load TPC-H Dataset, as shown in Figure 1.
After a short time, the database will be ready. The database name will start with s2_dataset_tpch_ and be visible under the Workspace Overview tab. We'll make a note of the entire database name.
As described in a previous article, we'll follow the instructions to create a Notebook.
Fill out the Notebook
First, we'll install LangChain:
!pip install langchain langchain-community langchain-openai --quiet
and add some imports:
from langchain_community.agent_toolkits.sql.base import create_sql_agent
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
Next, we'll set our OpenAI API Key:
import os
import getpass
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
We'll ensure that the workspace and database are selected at the top of the notebook and then create the agent:
db = SQLDatabase.from_uri(connection_url, include_tables = ["customer", "nation"])
llm = ChatOpenAI(model = "gpt-4o-mini", temperature = 0, verbose = False)
toolkit = SQLDatabaseToolkit(db = db, llm = llm)
agent_executor = create_sql_agent(
llm = llm,
toolkit = toolkit,
max_iterations = 15,
max_execution_time = 60,
top_k = 3,
verbose = False
)
Some code to handle errors:
error_string = "Could not parse LLM output:"
def run_agent_query(query, agent_executor, error_string):
try:
result = agent_executor.invoke(query, return_only_outputs = True)["output"]
except Exception as e:
error_message = str(e)
# Check if the error message contains the specific string
if error_string in error_message:
# Extract the part after the specific string and strip backticks
result = error_message.split(error_string)[1].strip().strip('`')
else:
result = f"Error occurred: {error_message}"
return result
Finally, we'll run the agent, as follows:
query = (
"Using the customer and nation tables, write a SingleStore query\n"
"that shows the highest paying customers per country,\n"
"include how much they've spent, use the nation name."
)
result = run_agent_query(query, agent_executor, error_string)
print(result)
Example output:
The highest paying customers per country are:
1. **Customer**: Customer#000508503 from **Country**: UNITED STATES with an amount spent of 9999.99.
2. **Customer**: Customer#000061453 from **Country**: MOROCCO with an amount spent of 9999.99.
3. **Customer**: Customer#001123705 from **Country**: MOZAMBIQUE with an amount spent of 9999.98.
Summary
Using the SQLDatabaseToolkit we can interact with a SingleStoreDB database using English. This provides an alternative to using SQL.
Acknowledgements
I thank my colleague Justin Kuntz for the demo code, which was adapted for this article.
Top comments (0)