""" nl2sql.py — natural-language -> SQL over a banking database, done SAFELY. This is one of the most-requested banking GenAI features ("ask questions about your transactions in plain language") and the JD's mandatory "Python and SQL". It is also a security minefield, so the lab's real lesson is the GUARD, not the generation. The pattern (Knowledge Modules 01 §8, 03 §12, 09 §9): NL question --LLM--> candidate SQL --GUARD--> execute (read-only, per-tenant views) --> rows --LLM--> answer Four safety layers, all enforced by the APP, never trusted to the model: 1. READ-ONLY — only SELECT is allowed; any DML/DDL is rejected. 2. ALLOWLIST — the SQL may only touch per-tenant VIEWS (my_accounts, my_transactions), never the base tables. 3. ROW-LEVEL SECURITY — those views are defined for the CALLER's customer_id (from the authenticated session), so a query physically cannot return another customer's data — even if the NL asks for "all customers". 4. PARAMETERIZATION — user-supplied values (a merchant name) are bound, never string-concatenated, so SQL injection is impossible. The GENERATOR is a deterministic template stand-in for an Azure OpenAI call; swap in AzureOpenAISqlGenerator and the guard/executor are unchanged (you NEVER trust generated SQL, regardless of which model wrote it). """ from __future__ import annotations import re import sqlite3 class SqlSafetyError(Exception): pass # --- the per-tenant views the model is allowed to query --------------------- ALLOWED_TABLES = {"my_accounts", "my_transactions"} _FORBIDDEN = re.compile( r"\b(insert|update|delete|drop|alter|attach|detach|pragma|create|replace|" r"truncate|grant|revoke|vacuum|reindex)\b", re.IGNORECASE) _TABLE_REF = re.compile(r"\b(?:from|join)\s+([a-zA-Z_][\w]*)", re.IGNORECASE) def guard(sql: str) -> str: """Validate generated SQL before it runs. Raises SqlSafetyError on anything that isn't a single read-only SELECT over the allowlisted per-tenant views.""" s = sql.strip().rstrip(";").strip() if ";" in s: raise SqlSafetyError("multiple statements are not allowed") if not s.lower().startswith("select"): raise SqlSafetyError("only SELECT queries are allowed (read-only)") if _FORBIDDEN.search(s): raise SqlSafetyError("forbidden keyword (DML/DDL) detected") tables = {t.lower() for t in _TABLE_REF.findall(s)} illegal = tables - ALLOWED_TABLES if illegal: raise SqlSafetyError(f"query references non-allowlisted tables: {sorted(illegal)}") return s class SqlExecutor: """Runs guarded SELECTs against per-tenant views bound to the CALLER. The customer_id comes from the authenticated session (trusted, cast to int) — it is NEVER taken from the natural-language question or the model.""" def __init__(self, conn: sqlite3.Connection) -> None: self.conn = conn def run(self, sql: str, params: dict, customer_id: int) -> list[dict]: safe = guard(sql) # layer 1+2 (read-only + allowlist) cid = int(customer_id) # layer 3 (RLS) — trusted int only cur = self.conn.cursor() cur.execute("DROP VIEW IF EXISTS my_accounts") cur.execute("DROP VIEW IF EXISTS my_transactions") cur.execute(f"CREATE TEMP VIEW my_accounts AS " f"SELECT * FROM accounts WHERE customer_id = {cid}") cur.execute(f"CREATE TEMP VIEW my_transactions AS " f"SELECT t.* FROM transactions t JOIN accounts a " f"ON t.account_id = a.id WHERE a.customer_id = {cid}") rows = cur.execute(safe, params).fetchall() # layer 4 (params bound) return [dict(r) for r in rows] # --- the NL -> SQL generator (deterministic stand-in for Azure OpenAI) ------- class TemplateSqlGenerator: """Maps a banking question to (sql, params) using parameterized templates over the per-tenant views. A real implementation prompts gpt-4o with the schema and the question and asks for SELECT-only SQL (structured output); the guard/executor then enforce safety regardless of what it produces.""" def generate(self, question: str) -> tuple[str, dict, str]: q = question.lower() if "balance" in q or "how much do i have" in q or "my accounts" in q: return ("SELECT type, balance, currency FROM my_accounts ORDER BY type", {}, "balances") if "category" in q or "breakdown" in q or "what did i spend on" in q: return ("SELECT category, ROUND(SUM(-amount),2) AS spent " "FROM my_transactions WHERE amount < 0 " "GROUP BY category ORDER BY spent DESC", {}, "spend_by_category") if "largest" in q or "biggest" in q: return ("SELECT date, description, amount FROM my_transactions " "ORDER BY ABS(amount) DESC LIMIT 1", {}, "largest_txn") # word-boundary match: avoid 'count' matching the substring in 'accounts'. if "how many" in q or "number of" in q or re.search(r"\bcount\b", q): return ("SELECT COUNT(*) AS n FROM my_transactions", {}, "count") if "income" in q or "salary" in q or "paid in" in q: return ("SELECT ROUND(SUM(amount),2) AS total_income " "FROM my_transactions WHERE amount > 0", {}, "income") merchant = _extract_merchant(question) if merchant: return ("SELECT date, description, amount FROM my_transactions " "WHERE description LIKE :merchant ORDER BY date", {"merchant": f"%{merchant}%"}, "by_merchant") if "spend" in q or "spent" in q or "spending" in q: return ("SELECT ROUND(SUM(-amount),2) AS total_spent " "FROM my_transactions WHERE amount < 0", {}, "total_spend") return ("", {}, "unknown") def _extract_merchant(question: str) -> str | None: """Pull a merchant token after 'at'/'from' or inside quotes. Whatever we extract is BOUND as a parameter, so a malicious value can't inject SQL.""" m = re.search(r"['\"]([^'\"]+)['\"]", question) if m: return m.group(1) m = re.search(r"\b(?:at|from|to)\s+([A-Za-z][\w&'-]+)", question) return m.group(1) if m else None class AzureOpenAISqlGenerator: """PRODUCTION generator (drop-in for TemplateSqlGenerator): SYSTEM = ("Generate a single read-only SQLite SELECT over ONLY these " "views: my_accounts(type,balance,currency), " "my_transactions(date,description,amount,category). " "Never reference other tables. Return JSON {sql, params}.") out = client.chat.completions.create(model="gpt-4o", temperature=0, response_format={"type":"json_object"}, messages=[{"role":"system","content":SYSTEM}, {"role":"user","content":question}]) # parse out['sql'], out['params'] — then GUARD + execute exactly as here. """ # --- the tool the agent calls ----------------------------------------------- class NL2SQLTool: """`answer(question, customer_id)` -> grounded natural-language answer + the exact SQL and rows (for audit). The answer is phrased FROM the returned rows, so the model never invents a number — same principle as every money fact.""" def __init__(self, conn: sqlite3.Connection, generator=None) -> None: self.executor = SqlExecutor(conn) self.generator = generator or TemplateSqlGenerator() def answer(self, question: str, customer_id: int) -> dict: sql, params, intent = self.generator.generate(question) if not sql: return {"answer": "I can't turn that into a query I'm allowed to run. " "Let me connect you with a human.", "sql": None, "rows": [], "intent": intent} rows = self.executor.run(sql, params, customer_id) return {"answer": _phrase(intent, rows), "sql": sql, "params": params, "rows": rows, "intent": intent} def _phrase(intent: str, rows: list[dict]) -> str: """Deterministic grounded phrasing from the real rows (a real LLM would do this nicer, but still only from the rows).""" if not rows: return "I didn't find any matching records for your account." if intent == "balances": return "Your accounts: " + "; ".join( f"{r['type']} {r['balance']:.2f} {r['currency']}" for r in rows) + "." if intent == "total_spend": return f"You spent {rows[0]['total_spent']:.2f} AED in total." if intent == "income": return f"Your total income is {rows[0]['total_income']:.2f} AED." if intent == "count": return f"You have {rows[0]['n']} transactions." if intent == "spend_by_category": return "Spending by category: " + ", ".join( f"{r['category']} {r['spent']:.2f}" for r in rows) + "." if intent == "largest_txn": r = rows[0] return f"Your largest transaction was {r['description']} for {r['amount']:.2f} AED on {r['date']}." if intent == "by_merchant": return "Matching transactions: " + ", ".join( f"{r['date']} {r['description']} {r['amount']:.2f}" for r in rows) + "." return str(rows)