nlp_sql.py

import os
import time
from typing import Dict, Any, Optional

import openai
import sqlglot
from sqlglot import parse_one, errors

openai.api_key = os.environ.get("OPENAI_API_KEY")
OPENAI_MODEL = os.environ.get("OPENAI_MODEL", "gpt-4")


def _call_openai_system(prompt: str, max_tokens: int = 512) -> str:
    resp = openai.ChatCompletion.create(
        model=OPENAI_MODEL,
        messages=[
            {"role": "system", "content": "You are an expert SQL generator. Return only the SQL query in a single code block without extra explanation."},
            {"role": "user", "content": prompt},
        ],
        max_tokens=max_tokens,
        temperature=0,
    )
    return resp["choices"][0]["message"]["content"].strip()


def clean_response_sql(text: str) -> str:
    # strip markdown fences
    if text.startswith("```"):
        # remove code fence markers
        parts = text.split("\n")
        if parts[0].startswith("```"):
            parts = parts[1:]
        if parts and parts[-1].startswith("```"):
            parts = parts[:-1]
        return "\n".join(parts).strip()
    return text


def validate_sql(sql: str, dialect: str = "default") -> Optional[str]:
    try:
        # sqlglot supports dialect names like 'postgres', 'mysql', 'sqlite'
        parse_one(sql, read=dialect if dialect != "default" else None)
        return None
    except Exception as e:
        return str(e)


def is_select_only(sql: str, dialect: str = "default") -> bool:
    try:
        expr = parse_one(sql, read=dialect if dialect != "default" else None)
        return expr is not None and expr.args.get("from") is not None and expr.__class__.__name__.lower() == "select"
    except Exception:
        return False


def generate_sql_from_nl(nl: str, semantic: Dict[str, Any], dialect: str = "default", max_retries: int = 5) -> Dict[str, Any]:
    prompt = f"Given the database schema and samples: {semantic}\n\nWrite a {dialect} SQL SELECT query that answers the user request: {nl}\nReturn only the SQL statement. No explanation."

    last_error = None
    for attempt in range(1, max_retries + 1):
        raw = _call_openai_system(prompt)
        sql = clean_response_sql(raw)
        err = validate_sql(sql, dialect=dialect)
        if err is None and is_select_only(sql, dialect=dialect):
            return {"sql": sql, "attempts": attempt, "error": None}
        last_error = err or "Not a SELECT statement"
        # ask model to correct with error
        prompt = f"Previous SQL:\n{sql}\n\nThe validator returned error: {last_error}. Fix the SQL and return only corrected SQL." 
        time.sleep(0.4)

    return {"sql": sql, "attempts": max_retries, "error": last_error}