SQL Agent with Cohere and LangChain (i-5O Case Study) ↗

This tutorial demonstrates how to create a SQL agent using Cohere and LangChain. The agent can translate natural language queries coming from users into SQL, and execute them against a database. This powerful combination allows for intuitive interaction with databases without requiring direct SQL knowledge.

Key topics covered:

1. Setting up the necessary libraries and environment
2. Connecting to a SQLite database
3. Configuring the LangChain SQL Toolkit
4. Creating a custom prompt template with few-shot examples
5. Building and running the SQL agent
6. Adding memory to the agent to keep track of historical messages

By the end of this tutorial, you’ll have a functional SQL agent that can answer questions about your data using natural language.

This tutorial uses a mocked up data of a manufacturing environment where a product item’s production is tracked across multiple stations, allowing for analysis of production efficiency, station performance, and individual item progress through the manufacturing process. This is modelled after a real customer use case.

The database contains two tables:

• The product_tracking table records the movement of items through different zones in manufacturing stations, including start and end times, station names, and product IDs.
• The status table logs the operational status of stations, including timestamps, station names, and whether they are productive or in downtime.

Import the required libraries#

First, let’s import the necessary libraries for creating a SQL agent using Cohere and LangChain. These libraries enable natural language interaction with databases and provide tools for building AI-powered agents.

import os

os.environ["COHERE_API_KEY"] = "<cohere-api-key>"

! pip install langchain-core langchain-cohere langchain-community faiss-cpu -qq

from langchain_cohere import create_sql_agent
from langchain_cohere.chat_models import ChatCohere
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.vectorstores import FAISS
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_cohere import CohereEmbeddings
from datetime import datetime

Load the database#

Next, we load the database for our manufacturing data.

Download the sql files from the link below to create the database.#

We create an in-memory SQLite database using SQL scripts for the product_tracking and status tables. You can get the SQL tables here.

We then create a SQLDatabase instance, which will be used by our LangChain tools and agents to interact with the data.

import sqlite3

from langchain_community.utilities.sql_database import SQLDatabase
from sqlalchemy import create_engine
from sqlalchemy.pool import StaticPool

def get_engine_for_manufacturing_db():
    """Create an in-memory database with the manufacturing data tables."""
    connection = sqlite3.connect(":memory:", check_same_thread=False)

    # Read and execute the SQL scripts
    for sql_file in ['product_tracking.sql', 'status.sql']:
        with open(sql_file, 'r') as file:
            sql_script = file.read()
            connection.executescript(sql_script)

    return create_engine(
        "sqlite://",
        creator=lambda: connection,
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )

engine = get_engine_for_manufacturing_db()

# Create the SQLDatabase instance
db = SQLDatabase(engine)

# Now you can use this db instance with your LangChain tools and agents

# Test the connection
db.run("SELECT * FROM status LIMIT 5;")

"[('2024-05-09 19:28:00', 'Canada/Toronto', '2024-05-09', '19', '28', 'stn3', 'downtime'), ('2024-04-21 06:57:00', 'Canada/Toronto', '2024-04-21', '6', '57', 'stn3', 'productive'), ('2024-04-11 23:52:00', 'Canada/Toronto', '2024-04-11', '23', '52', 'stn4', 'productive'), ('2024-04-03 21:52:00', 'Canada/Toronto', '2024-04-03', '21', '52', 'stn2', 'downtime'), ('2024-04-30 05:01:00', 'Canada/Toronto', '2024-04-30', '5', '1', 'stn4', 'productive')]"

# Test the connection
db.run("SELECT * FROM product_tracking LIMIT 5;")

"[('2024-05-27 17:22:00', '2024-05-27 17:57:00', 'Canada/Toronto', '2024-05-27', '17', 'stn2', 'wip', '187', '35'), ('2024-04-26 15:56:00', '2024-04-26 17:56:00', 'Canada/Toronto', '2024-04-26', '15', 'stn4', 'wip', '299', '120'), ('2024-04-12 04:36:00', '2024-04-12 05:12:00', 'Canada/Toronto', '2024-04-12', '4', 'stn3', 'wip', '60', '36'), ('2024-04-19 15:15:00', '2024-04-19 15:22:00', 'Canada/Toronto', '2024-04-19', '15', 'stn4', 'wait', '227', '7'), ('2024-04-24 19:10:00', '2024-04-24 21:07:00', 'Canada/Toronto', '2024-04-24', '19', 'stn4', 'wait', '169', '117')]"

Setup the LangChain SQL Toolkit#

Next, we initialize the LangChain SQL Toolkit and initialize the language model to use Cohere’s LLM. This prepares the necessary components for querying the SQL database using natural language.

## Define model to use
import os

MODEL = "command-a-03-2025"

llm = ChatCohere(
    model=MODEL,
    temperature=0.1,
    verbose=True
)


toolkit = SQLDatabaseToolkit(db=db, llm=llm)
context = toolkit.get_context()
tools = toolkit.get_tools()

print("**List of pre-defined Langchain Tools**")
print([tool.name for tool in tools])

**List of pre-defined Langchain Tools**
['sql_db_query', 'sql_db_schema', 'sql_db_list_tables', 'sql_db_query_checker']

Create a prompt template#

Next, we create a prompt template. In this section, we will introduce a simple system message, and then also show how we can improve the prompt by introducing few shot prompting examples in the later sections. The system message is used to communicate instructions or provide context to the model at the beginning of a conversation.

In this case, we provide the model with context on what sql dialect it should use, how many samples to query among other instructions.

from langchain_core.prompts import (
    PromptTemplate,
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    MessagesPlaceholder
)

system_message = """You are an agent designed to interact with a SQL database.
You are an expert at answering questions about manufacturing data.
Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Always start with checking the schema of the available tables.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most {top_k} results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the given tools. Only use the information returned by the tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

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

The current date is {date}.

For questions regarding productive time, downtime, productive or productivity, use minutes as units.

For questions regarding productive time, downtime, productive or productivity use the status table.

For questions regarding processing time and average processing time, use minutes as units.

For questions regarding bottlenecks, processing time and average processing time use the product_tracking table.

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

system_prompt = PromptTemplate.from_template(system_message)

full_prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate(prompt=system_prompt),
        MessagesPlaceholder(variable_name='chat_history', optional=True),
        ("human", "{input}"),
        MessagesPlaceholder("agent_scratchpad"),
    ]
)

prompt_val = full_prompt.invoke({
        "input": "What was the productive time for all stations today?",
        "top_k": 5,
        "dialect": "SQLite",
        "date":datetime.now(),
        "agent_scratchpad": [],
    })
print(prompt_val.to_string())

System: You are an agent designed to interact with a SQL database.
You are an expert at answering questions about manufacturing data.
Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
Always start with checking the schema of the available tables.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the given tools. Only use the information returned by the tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

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

The current date is 2025-03-13 09:21:55.403450.

For questions regarding productive time, downtime, productive or productivity, use minutes as units.

For questions regarding productive time, downtime, productive or productivity use the status table.

For questions regarding processing time and average processing time, use minutes as units.

For questions regarding bottlenecks, processing time and average processing time use the product_tracking table.

If the question does not seem related to the database, just return "I don't know" as the answer.
Human: What was the productive time for all stations today?

Create a few-shot prompt template#

In the above step, we’ve created a simple system prompt. Now, let us see how we can create a better few shot prompt template in this section. Few-shot examples are used to provide the model with context and improve its performance on specific tasks. In this case, we’ll prepare examples of natural language queries and their corresponding SQL queries to help the model generate accurate SQL statements for our database.

In this example, we use SemanticSimilarityExampleSelector to select the top k examples that are most similar to an input query out of all the examples available.

examples = [
    {
        "input": "What was the average processing time for all stations on April 3rd 2024?",
        "query": "SELECT station_name, AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND zone = 'wip' GROUP BY station_name ORDER BY station_name;",
    },
    {
        "input": "What was the average processing time for all stations on April 3rd 2024 between 4pm and 6pm?",
        "query": "SELECT station_name, AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND CAST(hour AS INTEGER) BETWEEN 16 AND 18 AND zone = 'wip' GROUP BY station_name ORDER BY station_name;",
    },
    {
        "input": "What was the average processing time for stn4 on April 3rd 2024?",
        "query": "SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND station_name = 'stn4' AND zone = 'wip';",
    },
    {
        "input": "How much downtime did stn2 have on April 3rd 2024?",
        "query": "SELECT COUNT(*) AS downtime_count FROM status WHERE date = '2024-04-03' AND station_name = 'stn2' AND station_status = 'downtime';",
    },
    {
        "input": "What were the productive time and downtime numbers for all stations on April 3rd 2024?",
        "query": "SELECT station_name, station_status, COUNT(*) as total_time FROM status WHERE date = '2024-04-03' GROUP BY station_name, station_status;",
    },
    {
        "input": "What was the bottleneck station on April 3rd 2024?",
        "query": "SELECT station_name, AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND zone = 'wip' GROUP BY station_name ORDER BY avg_processing_time DESC LIMIT 1;",
    },
    {
        "input": "Which percentage of the time was stn5 down in the last week of May?",
        "query": "SELECT SUM(CASE WHEN station_status = 'downtime' THEN 1 ELSE 0 END) * 100.0 / COUNT(*) AS percentage_downtime FROM status WHERE station_name = 'stn5' AND date >= '2024-05-25' AND date <= '2024-05-31';",
    },
]

example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    CohereEmbeddings(
        cohere_api_key=os.getenv("COHERE_API_KEY"), model="embed-v4.0"
    ),
    FAISS,
    k=5,
    input_keys=["input"],
)

from langchain_core.prompts import (
    ChatPromptTemplate,
    FewShotPromptTemplate,
    MessagesPlaceholder,
    PromptTemplate,
    SystemMessagePromptTemplate,
)

system_prefix = """You are an agent designed to interact with a SQL database.
You are an expert at answering questions about manufacturing data.
Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Always start with checking the schema of the available tables.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most {top_k} results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the given tools. Only use the information returned by the tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

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

The current date is {date}.

For questions regarding productive time, downtime, productive or productivity, use minutes as units.

For questions regarding productive time, downtime, productive or productivity use the status table.

For questions regarding processing time and average processing time, use minutes as units.

For questions regarding bottlenecks, processing time and average processing time use the product_tracking table.

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

Here are some examples of user inputs and their corresponding SQL queries:
"""

few_shot_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=PromptTemplate.from_template(
        "User input: {input}\nSQL query: {query}"
    ),
    input_variables=["input", "dialect", "top_k","date"],
    prefix=system_prefix,
    suffix="",
)

full_prompt = ChatPromptTemplate.from_messages(
    [
        # In the previous section, this was system_prompt instead without the few shot examples.
        # We can use either prompting style as required
        SystemMessagePromptTemplate(prompt=few_shot_prompt),
        ("human", "{input}"),
        MessagesPlaceholder("agent_scratchpad"),
    ]
)

# Example formatted prompt
prompt_val = full_prompt.invoke(
    {
        "input": "What was the productive time for all stations today?",
        "top_k": 5,
        "dialect": "SQLite",
        "date":datetime.now(),
        "agent_scratchpad": [],
    }
)
print(prompt_val.to_string())

System: You are an agent designed to interact with a SQL database.
You are an expert at answering questions about manufacturing data.
Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer.
Always start with checking the schema of the available tables.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 5 results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the given tools. Only use the information returned by the tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

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

The current date is 2025-03-13 09:22:22.275098.

For questions regarding productive time, downtime, productive or productivity, use minutes as units.

For questions regarding productive time, downtime, productive or productivity use the status table.

For questions regarding processing time and average processing time, use minutes as units.

For questions regarding bottlenecks, processing time and average processing time use the product_tracking table.

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

Here are some examples of user inputs and their corresponding SQL queries:


User input: What were the productive time and downtime numbers for all stations on April 3rd 2024?
SQL query: SELECT station_name, station_status, COUNT(*) as total_time FROM status WHERE date = '2024-04-03' GROUP BY station_name, station_status;

User input: What was the average processing time for all stations on April 3rd 2024?
SQL query: SELECT station_name, AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND zone = 'wip' GROUP BY station_name ORDER BY station_name;

User input: What was the average processing time for all stations on April 3rd 2024 between 4pm and 6pm?
SQL query: SELECT station_name, AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND CAST(hour AS INTEGER) BETWEEN 16 AND 18 AND zone = 'wip' GROUP BY station_name ORDER BY station_name;

User input: What was the bottleneck station on April 3rd 2024?
SQL query: SELECT station_name, AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND zone = 'wip' GROUP BY station_name ORDER BY avg_processing_time DESC LIMIT 1;

User input: What was the average processing time for stn4 on April 3rd 2024?
SQL query: SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_time FROM product_tracking WHERE date = '2024-04-03' AND station_name = 'stn4' AND zone = 'wip';
Human: What was the productive time for all stations today?

Create the agent#

Next, we create an instance of the SQL agent using the LangChain framework, specifically using create_sql_agent.

This agent will be capable of interpreting natural language queries, converting them into SQL queries, and executing them against our database. The agent uses the LLM we defined earlier, along with the SQL toolkit and the custom prompt we created.

agent = create_sql_agent(
   llm=llm,
   toolkit=toolkit,
   prompt=full_prompt,
   verbose=True
)

Run the agent#

Now, we can run the agent and test it with a few different queries.

# %%time
output=agent.invoke({
   "input": "Which stations had some downtime in the month of May 2024?",
    "date": datetime.now()
})
print(output['output'])

# Answer: stn2, stn3 and stn5 had some downtime in the month of May 2024.

> Entering new Cohere SQL Agent Executor chain...

Invoking: `sql_db_list_tablessql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will query the connected SQL database to find the stations that had some downtime in the month of May 2024.

sql_db_list_tablessql_db_list_tables is not a valid tool, try one of [sql_db_query, sql_db_schema, sql_db_list_tables, sql_db_query_checker].
Invoking: `sql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will query the connected SQL database to find the stations that had some downtime in the month of May 2024.

product_tracking, status
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking, status'}`
responded: I have found the following tables: product_tracking and status. I will now query the schema of these tables to understand their structure.


CREATE TABLE product_tracking (
    timestamp_start TEXT, 
    timestamp_end TEXT, 
    timezone TEXT, 
    date TEXT, 
    hour TEXT, 
    station_name TEXT, 
    zone TEXT, 
    product_id TEXT, 
    duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/


CREATE TABLE status (
    timestamp_event TEXT, 
    timezone TEXT, 
    date TEXT, 
    hour TEXT, 
    minute TEXT, 
    station_name TEXT, 
    station_status TEXT
)

/*
3 rows from status table:
timestamp_event	timezone	date	hour	minute	station_name	station_status
2024-05-09 19:28:00	Canada/Toronto	2024-05-09	19	28	stn3	downtime
2024-04-21 06:57:00	Canada/Toronto	2024-04-21	6	57	stn3	productive
2024-04-11 23:52:00	Canada/Toronto	2024-04-11	23	52	stn4	productive
*/
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking, status'}`
responded: I have found the following tables: product_tracking and status. I will now query the schema of these tables to understand their structure.


CREATE TABLE product_tracking (
    timestamp_start TEXT, 
    timestamp_end TEXT, 
    timezone TEXT, 
    date TEXT, 
    hour TEXT, 
    station_name TEXT, 
    zone TEXT, 
    product_id TEXT, 
    duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/


CREATE TABLE status (
    timestamp_event TEXT, 
    timezone TEXT, 
    date TEXT, 
    hour TEXT, 
    minute TEXT, 
    station_name TEXT, 
    station_status TEXT
)

/*
3 rows from status table:
timestamp_event	timezone	date	hour	minute	station_name	station_status
2024-05-09 19:28:00	Canada/Toronto	2024-05-09	19	28	stn3	downtime
2024-04-21 06:57:00	Canada/Toronto	2024-04-21	6	57	stn3	productive
2024-04-11 23:52:00	Canada/Toronto	2024-04-11	23	52	stn4	productive
*/
Invoking: `sql_db_query_checker` with `{'query': "SELECT DISTINCT station_name\r\nFROM status\r\nWHERE station_status = 'downtime'\r\n  AND SUBSTR(date, 1, 7) = '2024-05';"}`
responded: I have found that the status table contains the relevant information, including the station_status column which contains the values 'productive' and 'downtime'. I will now query the database to find the stations that had some downtime in the month of May 2024.

```sql
SELECT DISTINCT station_name
FROM status
WHERE station_status = 'downtime'
    AND SUBSTR(date, 1, 7) = '2024-05';
```
Invoking: `sql_db_query_checker` with `{'query': "SELECT DISTINCT station_name\r\nFROM status\r\nWHERE station_status = 'downtime'\r\n  AND SUBSTR(date, 1, 7) = '2024-05';"}`
responded: I have found that the status table contains the relevant information, including the station_status column which contains the values 'productive' and 'downtime'. I will now query the database to find the stations that had some downtime in the month of May 2024.

```sql
SELECT DISTINCT station_name
FROM status
WHERE station_status = 'downtime'
    AND SUBSTR(date, 1, 7) = '2024-05';
```
Invoking: `sql_db_query_checker` with `{'query': "SELECT DISTINCT station_name\r\nFROM status\r\nWHERE station_status = 'downtime'\r\n  AND SUBSTR(date, 1, 7) = '2024-05';"}`
responded: I have found that the status table contains the relevant information, including the station_status column which contains the values 'productive' and 'downtime'. I will now query the database to find the stations that had some downtime in the month of May 2024.

```sql
SELECT DISTINCT station_name
FROM status
WHERE station_status = 'downtime'
    AND SUBSTR(date, 1, 7) = '2024-05';
```
Invoking: `sql_db_query_checker` with `{'query': "SELECT DISTINCT station_name\r\nFROM status\r\nWHERE station_status = 'downtime'\r\n  AND SUBSTR(date, 1, 7) = '2024-05';"}`
responded: I have found that the status table contains the relevant information, including the station_status column which contains the values 'productive' and 'downtime'. I will now query the database to find the stations that had some downtime in the month of May 2024.

```sql
SELECT DISTINCT station_name
FROM status
WHERE station_status = 'downtime'
    AND SUBSTR(date, 1, 7) = '2024-05';
```
Invoking: `sql_db_query` with `{'query': "SELECT DISTINCT station_name\r\nFROM status\r\nWHERE station_status = 'downtime'\r\n  AND SUBSTR(date, 1, 7) = '2024-05';"}`
responded: I have checked the query and it is correct. I will now execute the query to find the stations that had some downtime in the month of May 2024.

[('stn3',), ('stn5',), ('stn2',)]
Invoking: `sql_db_query` with `{'query': "SELECT DISTINCT station_name\r\nFROM status\r\nWHERE station_status = 'downtime'\r\n  AND SUBSTR(date, 1, 7) = '2024-05';"}`
responded: I have checked the query and it is correct. I will now execute the query to find the stations that had some downtime in the month of May 2024.

[('stn3',), ('stn5',), ('stn2',)]The stations that had some downtime in the month of May 2024 are:
- stn3
- stn5
- stn2

> Finished chain.
The stations that had some downtime in the month of May 2024 are:
- stn3
- stn5
- stn2

output=agent.invoke({
   "input": "What is the average processing duration at stn5 in the wip zone?",
"date": datetime.now()
})
print(output['output'])

# Answer: 39.17 minutes

> Entering new Cohere SQL Agent Executor chain...

Invoking: `sql_db_list_tablessql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will write a query to find the average processing duration at stn5 in the wip zone.

sql_db_list_tablessql_db_list_tables is not a valid tool, try one of [sql_db_query, sql_db_schema, sql_db_list_tables, sql_db_query_checker].
Invoking: `sql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will write a query to find the average processing duration at stn5 in the wip zone.

product_tracking, status
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking'}`
responded: I will use the product_tracking table to find the average processing duration at stn5 in the wip zone.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking'}`
responded: I will use the product_tracking table to find the average processing duration at stn5 in the wip zone.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/
Invoking: `sql_db_query_checker` with `{'query': "SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_duration\nFROM product_tracking\nWHERE station_name = 'stn5' AND zone = 'wip';"}`
responded: I will use the product_tracking table to find the average processing duration at stn5 in the wip zone. The relevant columns are station_name, zone and duration.

```sql
SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_duration
FROM product_tracking
WHERE station_name = 'stn5' AND zone = 'wip';
```
Invoking: `sql_db_query_checker` with `{'query': "SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_duration\nFROM product_tracking\nWHERE station_name = 'stn5' AND zone = 'wip';"}`
responded: I will use the product_tracking table to find the average processing duration at stn5 in the wip zone. The relevant columns are station_name, zone and duration.

```sql
SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_duration
FROM product_tracking
WHERE station_name = 'stn5' AND zone = 'wip';
```
Invoking: `sql_db_query` with `{'query': "SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_duration\nFROM product_tracking\nWHERE station_name = 'stn5' AND zone = 'wip';"}`
responded: I will now execute the query to find the average processing duration at stn5 in the wip zone.

[(39.166666666666664,)]
Invoking: `sql_db_query` with `{'query': "SELECT AVG(CAST(duration AS INTEGER)) AS avg_processing_duration\nFROM product_tracking\nWHERE station_name = 'stn5' AND zone = 'wip';"}`
responded: I will now execute the query to find the average processing duration at stn5 in the wip zone.

[(39.166666666666664,)]The average processing duration at stn5 in the wip zone is 39.17 minutes.

> Finished chain.
The average processing duration at stn5 in the wip zone is 39.17 minutes.

output=agent.invoke({
   "input": "Which station had the highest total duration in the wait zone?",
"date": datetime.now()
})
print(output['output'])

# Answer: stn4 - 251 minutes

> Entering new Cohere SQL Agent Executor chain...

Invoking: `sql_db_list_tablessql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will write a query to find the station with the highest total duration in the wait zone.

sql_db_list_tablessql_db_list_tables is not a valid tool, try one of [sql_db_query, sql_db_schema, sql_db_list_tables, sql_db_query_checker].
Invoking: `sql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will write a query to find the station with the highest total duration in the wait zone.

product_tracking, status
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking'}`
responded: I will use the product_tracking table to find the station with the highest total duration in the wait zone. I will group the results by station_name and zone, and filter for the zone 'wait'. I will then order the results by total duration in descending order and limit the results to 1.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking'}`
responded: I will use the product_tracking table to find the station with the highest total duration in the wait zone. I will group the results by station_name and zone, and filter for the zone 'wait'. I will then order the results by total duration in descending order and limit the results to 1.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/
Invoking: `sql_db_query_checker` with `{'query': "SELECT station_name, SUM(CAST(duration AS INTEGER)) AS total_duration\r\nFROM product_tracking\r\nWHERE zone = 'wait'\r\nGROUP BY station_name\r\nORDER BY total_duration DESC\r\nLIMIT 1;"}`
responded: I will use the product_tracking table to find the station with the highest total duration in the wait zone. I will group the results by station_name and zone, and filter for the zone 'wait'. I will then order the results by total duration in descending order and limit the results to 1.

```sql
SELECT station_name, SUM(CAST(duration AS INTEGER)) AS total_duration
FROM product_tracking
WHERE zone = 'wait'
GROUP BY station_name
ORDER BY total_duration DESC
LIMIT 1;
```
Invoking: `sql_db_query` with `{'query': "SELECT station_name, SUM(CAST(duration AS INTEGER)) AS total_duration\nFROM product_tracking\nWHERE zone = 'wait'\nGROUP BY station_name\nORDER BY total_duration DESC\nLIMIT 1;"}`
responded: I will now execute the query to find the station with the highest total duration in the wait zone.

[('stn4', 251)]The station with the highest total duration in the wait zone is stn4, with a total duration of 251 minutes.

> Finished chain.
The station with the highest total duration in the wait zone is stn4, with a total duration of 251 minutes.

Memory in the sql agent#

We may want the agent to hold memory of our previous messages so that we’re able to coherently engage with the agent to answer our queries. In this section, let’s take a look at how we can add memory to the agent so that we’re able to achieve this outcome!

from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage
from pydantic import BaseModel, Field
from typing import List

In the code snippets below, we create a class to store the chat history in memory. This can be customised to store the messages from a database or any other suitable data store.

class InMemoryHistory(BaseChatMessageHistory, BaseModel):
"""In memory implementation of chat message history."""

messages: List[BaseMessage] = Field(default_factory=list)

def add_messages(self, messages: List[BaseMessage]) -> None:
    """Add a list of messages to the store"""
    self.messages.extend(messages)

def clear(self) -> None:
    self.messages = []

In the below code snippet, we make use of the RunnableWithMessageHistory abstraction to wrap around the agent we’ve created above to provide the message history to the agent that we can now utilize by chatting with the agent_with_chat_history as shown below.

store = {}


def get_by_session_id(session_id: str):
if session_id not in store:
    store[session_id] = InMemoryHistory()
return store[session_id]


agent_with_chat_history = RunnableWithMessageHistory(
agent, get_by_session_id, history_messages_key="chat_history"
)

output = agent_with_chat_history.invoke(
{
    "input": "What station had the longest duration on 27th May 2024?",
    "date": datetime.now(),
},
config={"configurable": {"session_id": "foo"}},
)
print(output["output"])

# Answer: stn2, with duration of 35 mins.

> Entering new Cohere SQL Agent Executor chain...

Invoking: `sql_db_list_tablessql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will query the product_tracking table to find the station with the longest duration on 27th May 2024.

sql_db_list_tablessql_db_list_tables is not a valid tool, try one of [sql_db_query, sql_db_schema, sql_db_list_tables, sql_db_query_checker].
Invoking: `sql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will query the product_tracking table to find the station with the longest duration on 27th May 2024.

product_tracking, status
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking'}`
responded: I will query the product_tracking table to find the station with the longest duration on 27th May 2024.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking'}`
responded: I will query the product_tracking table to find the station with the longest duration on 27th May 2024.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/
Invoking: `sql_db_query_checker` with `{'query': "SELECT station_name, duration\r\nFROM product_tracking\r\nWHERE date = '2024-05-27'\r\nORDER BY duration DESC\r\nLIMIT 1;"}`
responded: I will query the product_tracking table to find the station with the longest duration on 27th May 2024. I will filter the data for the date '2024-05-27' and order the results by duration in descending order to find the station with the longest duration.

```sql
SELECT station_name, duration
FROM product_tracking
WHERE date = '2024-05-27'
ORDER BY duration DESC
LIMIT 1;
```
Invoking: `sql_db_query` with `{'query': "SELECT station_name, duration\nFROM product_tracking\nWHERE date = '2024-05-27'\nORDER BY duration DESC\nLIMIT 1;"}`
responded: I will now execute the SQL query to find the station with the longest duration on 27th May 2024.

[('stn2', '35')]The station with the longest duration on 27th May 2024 was stn2 with a duration of 35 minutes.

> Finished chain.
The station with the longest duration on 27th May 2024 was stn2 with a duration of 35 minutes.

output = agent_with_chat_history.invoke(
{
    "input": "Can you tell me when this station had downtime on 2024-04-03?",
    "date": datetime.now(),
},
config={"configurable": {"session_id": "foo"}},
)
print(output["output"])

# Answer: 21:52:00

> Entering new Cohere SQL Agent Executor chain...

Invoking: `sql_db_list_tablessql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will query the database to find out when the station had downtime on 2024-04-03.

sql_db_list_tablessql_db_list_tables is not a valid tool, try one of [sql_db_query, sql_db_schema, sql_db_list_tables, sql_db_query_checker].
Invoking: `sql_db_list_tables` with `{}`
responded: I will first check the schema of the available tables. Then, I will query the database to find out when the station had downtime on 2024-04-03.

product_tracking, status
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking, status'}`
responded: I have found that there are two tables: product_tracking and status. I will now check the schema of these tables.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/


CREATE TABLE status (
timestamp_event TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
minute TEXT, 
station_name TEXT, 
station_status TEXT
)

/*
3 rows from status table:
timestamp_event	timezone	date	hour	minute	station_name	station_status
2024-05-09 19:28:00	Canada/Toronto	2024-05-09	19	28	stn3	downtime
2024-04-21 06:57:00	Canada/Toronto	2024-04-21	6	57	stn3	productive
2024-04-11 23:52:00	Canada/Toronto	2024-04-11	23	52	stn4	productive
*/
Invoking: `sql_db_schema` with `{'table_names': 'product_tracking, status'}`
responded: I have found that there are two tables: product_tracking and status. I will now check the schema of these tables.


CREATE TABLE product_tracking (
timestamp_start TEXT, 
timestamp_end TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
station_name TEXT, 
zone TEXT, 
product_id TEXT, 
duration TEXT
)

/*
3 rows from product_tracking table:
timestamp_start	timestamp_end	timezone	date	hour	station_name	zone	product_id	duration
2024-05-27 17:22:00	2024-05-27 17:57:00	Canada/Toronto	2024-05-27	17	stn2	wip	187	35
2024-04-26 15:56:00	2024-04-26 17:56:00	Canada/Toronto	2024-04-26	15	stn4	wip	299	120
2024-04-12 04:36:00	2024-04-12 05:12:00	Canada/Toronto	2024-04-12	4	stn3	wip	60	36
*/


CREATE TABLE status (
timestamp_event TEXT, 
timezone TEXT, 
date TEXT, 
hour TEXT, 
minute TEXT, 
station_name TEXT, 
station_status TEXT
)

/*
3 rows from status table:
timestamp_event	timezone	date	hour	minute	station_name	station_status
2024-05-09 19:28:00	Canada/Toronto	2024-05-09	19	28	stn3	downtime
2024-04-21 06:57:00	Canada/Toronto	2024-04-21	6	57	stn3	productive
2024-04-11 23:52:00	Canada/Toronto	2024-04-11	23	52	stn4	productive
*/
Invoking: `sql_db_query` with `{'query': "SELECT timestamp_event, station_name\r\nFROM status\r\nWHERE date = '2024-04-03'\r\nAND station_status = 'downtime';"}`
responded: I have found that the status table contains information about downtime. I will now query the database to find out when the station had downtime on 2024-04-03.

[('2024-04-03 21:52:00', 'stn2')]
Invoking: `sql_db_query` with `{'query': "SELECT timestamp_event, station_name\r\nFROM status\r\nWHERE date = '2024-04-03'\r\nAND station_status = 'downtime';"}`
responded: I have found that the status table contains information about downtime. I will now query the database to find out when the station had downtime on 2024-04-03.

[('2024-04-03 21:52:00', 'stn2')]The station stn2 had downtime at 21:52 on 2024-04-03.

> Finished chain.
The station stn2 had downtime at 21:52 on 2024-04-03.

We can see from the above code snippets that the agent is automatically able to infer and query with respect to ‘stn2’ in the above question without us having to specify it explicitly. This allows us to have more coherent conversations with the agent.

Conclusion#

This tutorial demonstrated how to create a SQL agent using Cohere and LangChain. The agent can translate natural language queries coming from users into SQL, and execute them against a database. This powerful combination allows for intuitive interaction with databases without requiring direct SQL knowledge.