File size: 2,150 Bytes
981d8b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """OpenRouter NL->SQL: schema-aware prompt + robust SQL extraction."""
from __future__ import annotations
import os
import re
from ._openrouter import get_api_key, make_client
DEFAULT_MODEL = "nvidia/nemotron-3-nano-30b-a3b:free"
SYSTEM_PROMPT = (
"You are an expert data analyst who writes DuckDB SQL. Given a database schema "
"and a question, return ONE read-only SELECT query that answers it.\n"
"Rules:\n"
"1. SELECT only — never INSERT/UPDATE/DELETE/DROP/ALTER or any DDL.\n"
"2. Use ONLY the tables and columns listed in the schema.\n"
"3. Return ONLY the SQL — no prose, no explanation, no markdown fences."
)
_FENCE = re.compile(r"```(?:sql)?\s*(.*?)```", re.S | re.I)
def get_model() -> str:
return os.environ.get("OPENROUTER_MODEL", DEFAULT_MODEL).strip()
def build_messages(question: str, schema_text: str) -> list[dict]:
user = f"Schema:\n{schema_text}\n\nQuestion: {question}\n\nSQL:"
return [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user},
]
def extract_sql(text: str) -> str:
"""Pull a SQL statement out of a model response (handles ```sql fences and prose)."""
if not text:
return ""
m = _FENCE.search(text)
sql = (m.group(1) if m else text).strip()
# Trim leading prose: jump to the first WITH/SELECT keyword.
low = sql.lower()
idx = min(
(i for i in (low.find("with "), low.find("select ")) if i != -1),
default=-1,
)
if idx > 0:
sql = sql[idx:]
return sql.strip().rstrip(";").strip()
def generate_sql(
question: str,
schema_text: str,
api_key: str | None = None,
model: str | None = None,
temperature: float = 0.0,
) -> str:
api_key = api_key or get_api_key()
if not api_key:
raise RuntimeError("OPENROUTER_API_KEY is not set. See the README.")
client = make_client(api_key)
resp = client.chat.completions.create(
model=model or get_model(),
messages=build_messages(question, schema_text),
temperature=temperature,
)
return extract_sql(resp.choices[0].message.content or "")
|