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

```mdx
"[('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.

```python PYTHON
## 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])
```

```mdx
**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.

```python PYTHON
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)

```

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

```python PYTHON
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())
```

```mdx
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.

```python PYTHON
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';",
    },
]
```

```python PYTHON
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"],
)
```

```python PYTHON
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="",
)
```

```python PYTHON
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"),
    ]
)
```

```python PYTHON