sql_analyst.py

import json
import sqlite3
from langchain_ollama import ChatOllama
from langchain_core.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from dotenv import load_dotenv

load_dotenv()

print("DEBUG: Initializing LLM")
import os

# Use the secret set in HF Spaces
hf_api_token = os.environ.get("HUGGINGFACEHUB_ACCESS_TOKEN")
if hf_api_token is None:
    print("WARNING: HUGGINGFACEHUB_ACCESS_TOKEN not found in environment!")

print("DEBUG: Initializing LLM with HF token")
llm = HuggingFaceEndpoint(
    model="Qwen/Qwen3-Coder-Next",
    task="text-generation",
    huggingfacehub_api_token=hf_api_token  # pass the env variable here
)
model = ChatHuggingFace(llm=llm)


# ---------------------------------------------------------------------------- #
# Pilot Few-Shot SQL Examples (must be defined before use)
# ---------------------------------------------------------------------------- #
COMPLEX_EXAMPLES = """
-- Example A: Basic SELECT
User: "List all data from orders."
SQL:
SELECT * FROM orders;

-- Example B: Lookup by ID
User: "Get customer with ID 5."
SQL:
SELECT * FROM customers WHERE customer_id = 5;

-- Example C: Aggregation with Group
User: "Count number of products per category."
SQL:
SELECT category, COUNT(*) AS count_per_category
FROM products
GROUP BY category;

-- Example D: Update record
User: "Update customer ID 10's email."
SQL:
UPDATE customers SET email = 'new@example.com' WHERE customer_id = 10;

-- Example E: Delete specific records
User: "Delete cancelled orders before 2025."
SQL:
DELETE FROM orders WHERE status = 'cancelled' AND order_date < '2025-01-01';
"""

# ---------------------------------------------------------------------------- #
# Prompt Template (stores DDL, examples, and user query)
# ---------------------------------------------------------------------------- #
prompt_template = PromptTemplate(
    template="""
-- You are an expert SQL generator.
-- Generate valid SQL (SELECT, INSERT, UPDATE, DELETE) to answer the user's request.
-- Avoid hallucinating tables/columns outside the schema.
-- Do not output CREATE, DROP, ALTER, or other schema changes.
-- If the question cannot be answered with the given schema, output:
-- Cannot answer with current schema.

-- SQL Dialect: {sql_dialect}

-- Schema DDL --
{schema_ddl}

-- Examples --
{few_shot_examples}

-- User Request --
{user_query}

-- OUTPUT ONLY the SQL:
""",
    input_variables=["schema_ddl", "few_shot_examples", "user_query", "sql_dialect"]
)

# ---------------------------------------------------------------------------- #
# Main: Generate + Execute SQL
# ---------------------------------------------------------------------------- #
def generate_and_execute_sql(user_query: str, schema: dict, db_config: dict = None):
    print("\nDEBUG: generate_and_execute_sql called")
    print("DEBUG: User query:", user_query)
    print("DEBUG: DB config:", db_config)

    # Determine SQL dialect
    if db_config is None or db_config.get("type") == "sqlite":
        dialect = "SQLite"
    elif db_config["type"] == "mysql":
        dialect = "MySQL"
    elif db_config["type"] == "postgresql":
        dialect = "PostgreSQL"
    else:
        dialect = "SQL"

    print("DEBUG: Using dialect:", dialect)

    ddl = convert_schema_to_ddl(schema)

    # Build final prompt with the schema and examples
    final_prompt = prompt_template.format(
        schema_ddl=ddl,
        few_shot_examples=COMPLEX_EXAMPLES.strip(),
        user_query=user_query,
        sql_dialect=dialect
    )

    print("DEBUG: Prompt sent to model:\n", final_prompt)

    # Ask LLM for SQL
    try:
        response = model.invoke([
            ("system", "Generate only SQL as instructed."),
            ("human", final_prompt),
        ])
    except Exception as e:
        print("ERROR: LLM invocation error:", e)
        return "", f"LLM invocation error: {str(e)}"

    generated_sql = response.content.strip()
    print("DEBUG: Model generated SQL:\n", generated_sql)

    # If the LLM says it cannot answer:
    if generated_sql.lower().startswith("cannot answer"):
        print("DEBUG: Model cannot answer with current schema")
        return generated_sql, "No execution (cannot answer)"

    # Execute the generated SQL
    print("DEBUG: Executing SQL...")
    result = execute_query(generated_sql, db_config)
    print("DEBUG: Execution result:", result)

    return generated_sql, result

# ---------------------------------------------------------------------------- #
# Unified SQL Executor (with column names)
# ---------------------------------------------------------------------------- #
def execute_query(query: str, db_config: dict):
    print("\nDEBUG: execute_query called")
    print("DEBUG: Query to execute:", query)
    print("DEBUG: DB Config:", db_config)

    db_type = db_config.get("type")

    try:
        if db_type == "mysql":
            print("DEBUG: Connecting to MySQL")
            import mysql.connector
            conn = mysql.connector.connect(
                host=db_config["host"],
                port=int(db_config["port"]),
                user=db_config["user"],
                password=db_config["password"],
                database=db_config["database"],
            )

        elif db_type == "postgresql":
            print("DEBUG: Connecting to PostgreSQL")
            import psycopg2
            conn = psycopg2.connect(
                host=db_config["host"],
                port=int(db_config["port"]),
                user=db_config["user"],
                password=db_config["password"],
                dbname=db_config["database"],
            )
            print("DEBUG: PostgreSQL connection successful")

        else:
            print("DEBUG: Connecting to SQLite")
            conn = sqlite3.connect(db_config.get("path"))

        cursor = conn.cursor()
        print("DEBUG: Cursor acquired, executing now...")

        cursor.execute(query)

        # SELECT → fetch results + columns
        if query.strip().lower().startswith("select"):
            print("DEBUG: SELECT query — fetching rows")
            rows = cursor.fetchall()

            # Get column names from cursor.description
            columns = [col[0] for col in cursor.description]  # first element is column name :contentReference[oaicite:1]{index=1}

            print("DEBUG: Column names:", columns)
            print("DEBUG: SELECT result rows:", rows)

            conn.close()
            return {"columns": columns, "rows": rows}

        # For INSERT/UPDATE/DELETE
        print("DEBUG: Non‑SELECT — commit and return affected count")
        conn.commit()
        affected = cursor.rowcount
        print("DEBUG: Affected rows:", affected)
        conn.close()
        return {"affected_rows": affected}

    except Exception as e:
        print("ERROR: Exception in execute_query:", e)
        try:
            conn.close()
        except:
            pass

        return f"SQL execution error: {str(e)}"

# ---------------------------------------------------------------------------- #
# Schema to DDL
# ---------------------------------------------------------------------------- #
def convert_schema_to_ddl(schema: dict) -> str:
    print("DEBUG: Converting schema to DDL")
    ddl_lines = []
    for table, cols in schema.items():
        col_defs = []
        for col in cols:
            col_name = (
                col.get("name")
                or col.get("Field")
                or col.get("column_name")
                or col.get("COLUMN_NAME")
                or str(col)
            )
            col_type = (
                col.get("type")
                or col.get("Type")
                or col.get("data_type")
                or "TEXT"
            )
            col_defs.append(f"{col_name} {col_type}")
        ddl_lines.append(f"CREATE TABLE {table} ({', '.join(col_defs)});")
    ddl_text = "\n".join(ddl_lines)
    print("DEBUG: DDL:\n", ddl_text)
    return ddl_text