| """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() |
| |
| 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 "") |
|
|
