sqlcoder.py

"""
NaturalSQL — AI-Powered Natural Language to SQL Engine
======================================================
Uses Defog's SQLCoder-7b-2 model to convert plain-English questions
into PostgreSQL queries and execute them against a Cloud SQL database.

Usage:
    1. Copy `.env.template` to `.env` and fill in your credentials.
    2. Run:  python sqlcoder.py
"""

import os
import sys

import sqlparse
import torch
import pandas as pd
from dotenv import load_dotenv
from transformers import AutoTokenizer, AutoModelForCausalLM
from google.cloud.sql.connector import Connector

# ──────────────────────────────────────────────
# 1. Load environment variables
# ──────────────────────────────────────────────
load_dotenv()

CLOUD_SQL_INSTANCE = os.getenv("CLOUD_SQL_INSTANCE")
DB_DRIVER = os.getenv("DB_DRIVER", "pg8000")
DB_USER = os.getenv("DB_USER")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_NAME = os.getenv("DB_NAME")
MODEL_NAME = os.getenv("MODEL_NAME", "defog/sqlcoder-7b-2")

_REQUIRED_VARS = {
    "CLOUD_SQL_INSTANCE": CLOUD_SQL_INSTANCE,
    "DB_USER": DB_USER,
    "DB_PASSWORD": DB_PASSWORD,
    "DB_NAME": DB_NAME,
}
_missing = [k for k, v in _REQUIRED_VARS.items() if not v]
if _missing:
    sys.exit(
        f"[ERROR] Missing required environment variables: {', '.join(_missing)}\n"
        f"       Copy .env.template → .env and fill in the values."
    )

# ──────────────────────────────────────────────
# 2. Load the model & tokenizer
# ──────────────────────────────────────────────
print(f"🔄 Loading model: {MODEL_NAME} …")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

if torch.cuda.is_available():
    available_memory = torch.cuda.get_device_properties(0).total_memory
    if available_memory > 15e9:
        # ≥ 15 GB VRAM → full float16
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            trust_remote_code=True,
            torch_dtype=torch.float16,
            device_map="auto",
            use_cache=True,
        )
    else:
        # < 15 GB VRAM → 8-bit quantisation
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_NAME,
            trust_remote_code=True,
            load_in_8bit=True,
            device_map="auto",
            use_cache=True,
        )
    DEVICE = "cuda"
else:
    print("⚠️  No CUDA GPU detected — running on CPU (will be slow).")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
        torch_dtype=torch.float32,
        device_map="auto",
        use_cache=True,
    )
    DEVICE = "cpu"

print("✅ Model loaded successfully.\n")

# ──────────────────────────────────────────────
# 3. Prompt template (PostgreSQL schema)
# ──────────────────────────────────────────────
PROMPT_TEMPLATE = """\
### Task
Generate a SQL query to answer [QUESTION]{question}[/QUESTION]

### Instructions
- If you cannot answer the question with the available database schema, return 'I do not know.'
- Remember that revenue is calculated as the product's `price` multiplied by the `quantity` sold.
- Remember that cost is calculated as the `supply_price` from `product_suppliers` multiplied by the `quantity` sold.
- Ensure the SQL query is compatible with **PostgreSQL**. Avoid using features not supported by PostgreSQL, such as MySQL-specific functions. Instead, leverage PostgreSQL-compatible functions and syntax.

### Database Schema
This query will run on a database with the following schema:

CREATE TABLE products (
  product_id SERIAL PRIMARY KEY,
  name VARCHAR(50),
  price DECIMAL(10, 2),
  quantity INTEGER
);

CREATE TABLE customers (
  customer_id SERIAL PRIMARY KEY,
  name VARCHAR(50),
  address VARCHAR(100)
);

CREATE TABLE salespeople (
  salesperson_id SERIAL PRIMARY KEY,
  name VARCHAR(50),
  region VARCHAR(50)
);

CREATE TABLE sales (
  sale_id SERIAL PRIMARY KEY,
  product_id INTEGER,
  customer_id INTEGER,
  salesperson_id INTEGER,
  sale_date DATE,
  quantity INTEGER,
  FOREIGN KEY (product_id) REFERENCES products (product_id),
  FOREIGN KEY (customer_id) REFERENCES customers (customer_id),
  FOREIGN KEY (salesperson_id) REFERENCES salespeople (salesperson_id)
);

CREATE TABLE product_suppliers (
  supplier_id SERIAL PRIMARY KEY,
  product_id INTEGER,
  supply_price DECIMAL(10, 2),
  FOREIGN KEY (product_id) REFERENCES products (product_id)
);

CREATE TABLE product_categories (
  category_id SERIAL PRIMARY KEY,
  name VARCHAR(50)
);

CREATE TABLE orders (
  order_id SERIAL PRIMARY KEY,
  customer_id INTEGER,
  order_date DATE,
  total_amount DECIMAL(10, 2),
  FOREIGN KEY (customer_id) REFERENCES customers (customer_id)
);

CREATE TABLE inventory (
  inventory_id SERIAL PRIMARY KEY,
  product_id INTEGER,
  stock_level INTEGER,
  FOREIGN KEY (product_id) REFERENCES products (product_id)
);

CREATE TABLE payments (
  payment_id SERIAL PRIMARY KEY,
  order_id INTEGER,
  payment_date DATE,
  amount DECIMAL(10, 2),
  FOREIGN KEY (order_id) REFERENCES orders (order_id)
);

CREATE TABLE reviews (
  review_id SERIAL PRIMARY KEY,
  product_id INTEGER,
  customer_id INTEGER,
  rating INTEGER,
  review_date DATE,
  FOREIGN KEY (product_id) REFERENCES products (product_id),
  FOREIGN KEY (customer_id) REFERENCES customers (customer_id)
);

-- Relationships:
-- sales.product_id        → products.product_id
-- sales.customer_id       → customers.customer_id
-- sales.salesperson_id    → salespeople.salesperson_id
-- product_suppliers.product_id → products.product_id
-- orders.customer_id      → customers.customer_id
-- inventory.product_id   → products.product_id
-- payments.order_id       → orders.order_id
-- reviews.product_id      → products.product_id
-- reviews.customer_id     → customers.customer_id

### Answer
Given the database schema, here is the SQL query that answers [QUESTION]{question}[/QUESTION] using PostgreSQL:
[SQL]
"""

# ──────────────────────────────────────────────
# 4. SQL generation
# ──────────────────────────────────────────────
def generate_query(question: str) -> str:
    """Convert a natural-language question into a PostgreSQL query."""
    updated_prompt = PROMPT_TEMPLATE.format(question=question)
    inputs = tokenizer(updated_prompt, return_tensors="pt").to(DEVICE)

    generated_ids = model.generate(
        **inputs,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
        max_new_tokens=400,
        do_sample=False,
        num_beams=1,
    )
    outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

    return sqlparse.format(outputs[0].split("[SQL]")[-1], reindent=True)


# ──────────────────────────────────────────────
# 5. Database connection & execution
# ──────────────────────────────────────────────
connector = Connector()


def getconn():
    """Return a pg8000 connection via the Cloud SQL Python Connector."""
    return connector.connect(
        CLOUD_SQL_INSTANCE,
        DB_DRIVER,
        user=DB_USER,
        password=DB_PASSWORD,
        db=DB_NAME,
    )


def execute_sql_query(query: str) -> pd.DataFrame | None:
    """Execute a SQL query and return results as a DataFrame."""
    connection = getconn()
    try:
        cursor = connection.cursor()
        cursor.execute(query)
        columns = [desc[0] for desc in cursor.description]
        rows = cursor.fetchall()
        cursor.close()
        return pd.DataFrame(rows, columns=columns)
    except Exception as e:
        print(f"❌ Error executing query: {e}")
        return None
    finally:
        connection.close()


# ──────────────────────────────────────────────
# 6. Interactive loop
# ──────────────────────────────────────────────
def main():
    print("=" * 60)
    print("  NaturalSQL — Ask questions in plain English")
    print("  Type 'quit' or 'exit' to stop.")
    print("=" * 60)

    while True:
        question = input("\n❓ Your question: ").strip()
        if not question:
            continue
        if question.lower() in ("quit", "exit"):
            print("👋 Goodbye!")
            break

        print("\n🔄 Generating SQL …")
        sql = generate_query(question)
        print(f"\n📝 Generated SQL:\n{sql}")

        run = input("\nExecute this query? [Y/n]: ").strip().lower()
        if run in ("", "y", "yes"):
            print("\n⏳ Running query …")
            results = execute_sql_query(sql)
            if results is not None:
                print(f"\n✅ Results ({len(results)} rows):\n")
                print(results.to_string(index=False))
            else:
                print("⚠️  No results returned.")


if __name__ == "__main__":
    main()