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.

Figure 1. Load TPC-H Dataset.

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.

Quick tip: Using LangChain’s SQLDatabaseToolkit with SingleStoreDB