verbose live UI + tool-level SSE events + Groq default + regression tests

api/tools/text_to_sql.py

@@ -0,0 +1,70 @@
+"""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 []