71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
"""LLM-driven NL → SQL with sqlglot validation and one retry."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import re
|
|
|
|
import sqlglot
|
|
|
|
from api import langfuse_client as lf
|
|
from api.llm import chat
|
|
from api.prompts import load, render
|
|
from api.tools.schema import load_schema
|
|
from api.tools.types import T2SResult
|
|
|
|
logger = logging.getLogger("nvi.tools.text_to_sql")
|
|
|
|
|
|
def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SResult:
|
|
schema = load_schema()
|
|
tables = hint_tables or schema.table_names()
|
|
system = load("text_to_sql.system")
|
|
|
|
def _user(retry_hint: str = "") -> str:
|
|
return render(
|
|
"text_to_sql.user",
|
|
question=question,
|
|
schema_block=schema.render_tables(tables),
|
|
metrics_block=schema.render_metrics(),
|
|
retry_hint=retry_hint,
|
|
)
|
|
|
|
with lf.span(
|
|
"text_to_sql",
|
|
as_type="generation",
|
|
input={"question": question, "hint_tables": hint_tables},
|
|
) as span:
|
|
sql = _extract_sql(chat(system=system, user=_user()))
|
|
try:
|
|
_validate(sql)
|
|
except Exception as e:
|
|
logger.info("t2s first attempt invalid (%s); retrying once", e)
|
|
hint = f"\n\nPrevious attempt failed parsing: {e}. Fix and return only the SQL."
|
|
sql = _extract_sql(chat(system=system, user=_user(hint)))
|
|
_validate(sql)
|
|
|
|
result = T2SResult(sql=sql, used_tables=_extract_tables(sql))
|
|
span.update(output={"sql": sql, "used_tables": result.used_tables})
|
|
return result
|
|
|
|
|
|
def _extract_sql(text: str) -> str:
|
|
m = re.search(r"```sql\s*(.*?)```", text, re.DOTALL | re.IGNORECASE)
|
|
if m:
|
|
return m.group(1).strip().rstrip(";").strip()
|
|
return text.strip().rstrip(";").strip()
|
|
|
|
|
|
def _validate(sql: str) -> None:
|
|
parsed = sqlglot.parse_one(sql, dialect="postgres")
|
|
if parsed is None:
|
|
raise ValueError("sqlglot returned no parse tree")
|
|
|
|
|
|
def _extract_tables(sql: str) -> list[str]:
|
|
try:
|
|
parsed = sqlglot.parse_one(sql, dialect="postgres")
|
|
return sorted({t.name for t in parsed.find_all(sqlglot.exp.Table)})
|
|
except Exception:
|
|
return []
|