ex_1.py

from langchain.prompts import ChatPromptTemplate
from utils.db_tools import get_engine
from tools.table_schemas import get_table_schemas

query_generation = """
You are a SQL query expert capable of analyzing questions and generating error-free SQL queries to answer them.

## Main Tasks

1. Thoroughly understand the question before generating SQL queries.
2. Decide query type: Aggregate (using aggregation functions) or Normal query.
3. Write SQL queries to answer the question precisely.
4. Ensure the SQL queries are free of errors.

## Resources
- Schema: 
{SCHEMA}

- Categorical values:
{CATEGORICAL_VALUES}

### Step 1: Question Analysis

- Identify relevant columns from Schema.
- Recognize syntactically similar column names or values.
- Determine if numerical results are required.
- Identify necessary filters.
- Assess need for aggregation functions.
- Determine if joins or sub-queries are needed.

### Step 2: SQL Query Writing

Requirements to write correct SQL queries:
1. Select relevant columns from Step 1.
2. Use aliases, especially for aggregation functions.
3. Use sub-queries/CTEs and joins for complex queries.
4. Apply WHERE clause for filtering if needed.
5. Filter NULL values.
6. Avoid ORDER BY unless required.
7. Limit to 10 records (if user want to limit output results).
8. Ignore limit for aggregation functions.
9. Use LIKE with wildcards (%, _) instead of direct string comparison.
10. Use IN operator for multiple category comparisons.
11. Include grouping columns in SELECT clause
12. cast date columns to TEXT for better readability.
13. Avoid DML statements.
14. Sanitize input to prevent SQL injections.

### Step 3: Error Checking

- Check for SQL syntax errors.
- Check for logical errors.
- Verify query matches Step 1 actions.
- Correct any SQL syntax errors.

## Response Guidelines

**YOU MUST CALL THE CORRECT TOOL.**
<Tools>
    <Tool>db_query_tool</Tool>
</Tools>
"""

schemas = get_table_schemas(['payment', 'rental', 'inventory', 'film'], get_engine())
table_schema = ''

for table, schema in schemas.items():
    table_schema += f'Table: {table}\n'
    table_schema += schema
    table_schema += '\n\n'

categorical_desc = "rating column only can be one of, [PG, G, NC-17, PG-13, R]"

query_generation = ChatPromptTemplate.from_messages([
    ('system', query_generation.format(SCHEMA=table_schema,
                                       CATEGORICAL_VALUES=categorical_desc)),
    ('placeholder', '{messages}')
])