"""Read-only SQL tool. Three jobs: 1. **Introspect the schema once at startup** so the LLM has table / column types in its prompt without us hand-maintaining a catalog. 2. **Translate a natural-language intent into a single ``SELECT``** via ``ChatOpenAI`` with structured output. Refuses anything that isn't a read query. 3. **Execute the SQL safely.** Defense in depth: - The DB connection should already be a read-only role (``adaptive_rag_ro`` in our seed script). - Statement-level allowlist: only ``SELECT`` / ``WITH`` allowed. - ``statement_timeout`` set per-session so runaway queries die fast. - A ``LIMIT N`` is appended if the SQL doesn't already have one, so accidentally returning a million rows can't OOM the UI. This is intentionally a thin tool — it is *not* an agent. The dispatcher calls it once with a NL intent, and we either return rows or raise. """ from __future__ import annotations import logging import re from dataclasses import dataclass, field from typing import Any from langchain_openai import ChatOpenAI from pydantic import BaseModel, Field from sqlalchemy import create_engine, inspect, text from sqlalchemy.engine import Engine from sqlalchemy.exc import SQLAlchemyError from src.config import settings logger = logging.getLogger(__name__) # Single-statement, must start with SELECT or WITH (CTEs that resolve to a # SELECT). Reject anything that smells like a write or DDL even if it would # also be blocked by the read-only role — defense in depth, and surfaces # better error messages than letting Postgres reject mid-query. _SELECT_PATTERN = re.compile(r"^\s*(?:WITH\b|SELECT\b)", re.IGNORECASE) _FORBIDDEN_KEYWORDS = re.compile( r"\b(?:INSERT|UPDATE|DELETE|MERGE|TRUNCATE|DROP|ALTER|CREATE|" r"GRANT|REVOKE|COPY|VACUUM|ANALYZE|REINDEX|CLUSTER|" r"COMMENT\s+ON|SECURITY\s+DEFINER|DO\s+\$\$)\b", re.IGNORECASE, ) _LIMIT_PATTERN = re.compile(r"\blimit\s+\d+\b", re.IGNORECASE) _MULTI_STATEMENT = re.compile(r";\s*\S") SQL_SYSTEM_PROMPT = """You translate natural-language data questions into a single read-only PostgreSQL SELECT query. Hard rules: - Output exactly one statement. - The statement MUST start with SELECT or WITH. - NEVER write, modify or define schema (no INSERT, UPDATE, DELETE, CREATE, DROP, ALTER, etc.). The connection is read-only and will reject them anyway, but don't generate them. - Use only the tables and columns shown in the schema below. - When the question implies a time window like "last month" or "this year", prefer ``WHERE col >= NOW() - INTERVAL 'N units'`` over hardcoded dates. - Add an explicit ORDER BY when the question implies ranking ("top", "most", "biggest"). - If the question is ambiguous or cannot be answered from this schema, return a single SELECT that explains the gap, e.g. ``SELECT 'cannot answer: ' AS error;`` - Output only the SQL — no commentary, no markdown fences. """ SQL_USER_TEMPLATE = """Schema: {schema} Question: {intent} """ class _SqlOutput(BaseModel): """Structured output the SQL LLM is forced into.""" sql: str = Field( ..., description=( "A single PostgreSQL SELECT statement that answers the question. " "No trailing semicolon, no markdown fences." ), ) class SqlToolError(Exception): """Raised when the SQL tool can't produce or execute a safe query.""" @dataclass class SqlResult: """One execution of one query.""" intent: str sql: str columns: list[str] = field(default_factory=list) rows: list[dict[str, Any]] = field(default_factory=list) truncated: bool = False elapsed_ms: float = 0.0 @property def row_count(self) -> int: return len(self.rows) class SqlTool: """Schema-aware NL\u2192SQL tool with read-only execution.""" def __init__( self, database_url: str | None = None, *, statement_timeout_sec: int | None = None, row_limit: int | None = None, translator_model: str | None = None, ) -> None: url = database_url or settings.SQL_DATABASE_URL if not url: raise SqlToolError( "SQL_DATABASE_URL is not set. The SQL tool needs a connection " "string. Run scripts/seed_demo_data.py first or point at your " "own Postgres." ) self.database_url = url self.statement_timeout_sec = ( statement_timeout_sec or settings.SQL_QUERY_TIMEOUT_SEC ) self.row_limit = row_limit or settings.SQL_ROW_LIMIT try: self._engine: Engine = create_engine(url, future=True, pool_pre_ping=True) except SQLAlchemyError as exc: raise SqlToolError(f"Cannot create SQL engine: {exc}") from exc # Cache the schema description — it doesn't change at runtime and # we'd otherwise pay an introspection round-trip on every query. self._schema_text = self._describe_schema() self._llm = ChatOpenAI( model=translator_model or settings.SQL_MODEL, # Ignored by gpt-5.6 reasoning models (only temperature=1 is # supported); langchain-openai strips it automatically. temperature=0.0, # Pro-mode reasoning tokens count toward this budget — 400 was # enough for gpt-4.1-mini but truncates Luna Pro mid-reasoning. max_tokens=8000, reasoning=settings.sql_reasoning, ).with_structured_output(_SqlOutput) logger.info( f"SqlTool ready (db={self._safe_url()}, " f"timeout={self.statement_timeout_sec}s, row_limit={self.row_limit})" ) # ---- public API --------------------------------------------------- @property def schema_text(self) -> str: """Human-readable schema description (cached).""" return self._schema_text def schema_summary(self) -> str: """One-line-per-table summary, suitable for the router prompt.""" try: inspector = inspect(self._engine) lines: list[str] = [] for table in sorted(inspector.get_table_names(schema="public")): cols = [c["name"] for c in inspector.get_columns(table, schema="public")] preview = ", ".join(cols[:6]) if len(cols) > 6: preview += ", \u2026" lines.append(f"- {table}: {preview}") return "\n".join(lines) except SQLAlchemyError as exc: logger.warning(f"Schema summary failed: {exc}") return "(schema introspection failed)" def answer(self, intent: str) -> SqlResult: """End-to-end: NL intent -> SQL -> rows.""" sql = self.translate(intent) return self.execute(sql, intent=intent) def translate(self, intent: str) -> str: """Ask the LLM for a single SELECT statement matching ``intent``.""" intent = (intent or "").strip() if not intent: raise SqlToolError("Empty intent — nothing to translate.") messages = [ {"role": "system", "content": SQL_SYSTEM_PROMPT}, { "role": "user", "content": SQL_USER_TEMPLATE.format( schema=self._schema_text, intent=intent, ), }, ] from src.observability import get_callback_handler try: output: _SqlOutput = self._llm.invoke( messages, config={ "callbacks": get_callback_handler(), "run_name": "sql_tool.translate", "metadata": {"langfuse_tags": ["sql_tool", "nl2sql"]}, }, ) except Exception as exc: raise SqlToolError(f"LLM SQL translation failed: {exc}") from exc sql = self._clean(output.sql) self._validate(sql) return sql def execute(self, sql: str, *, intent: str = "") -> SqlResult: """Validate and run ``sql``, returning a :class:`SqlResult`.""" sql = self._clean(sql) self._validate(sql) sql_to_run = self._inject_limit(sql) import time t0 = time.perf_counter() try: with self._engine.connect() as conn: # Per-session timeout so runaway plans die fast. Postgres # accepts an integer string of milliseconds. conn.exec_driver_sql( f"SET statement_timeout = {self.statement_timeout_sec * 1000}" ) # Force read-only at the transaction level too. With a RO # role this is redundant, but it makes the intent explicit # and protects against misconfigured connection strings. conn.exec_driver_sql("SET TRANSACTION READ ONLY") result = conn.execute(text(sql_to_run)) rows = result.mappings().all() columns = list(result.keys()) except SQLAlchemyError as exc: raise SqlToolError(f"SQL execution failed: {exc}") from exc elapsed_ms = (time.perf_counter() - t0) * 1000 truncated = len(rows) >= self.row_limit and not _LIMIT_PATTERN.search(sql) materialised = [dict(r) for r in rows] logger.info( f"SqlTool.execute: {len(materialised)} rows ({elapsed_ms:.0f}ms) " f"truncated={truncated}" ) return SqlResult( intent=intent, sql=sql, columns=columns, rows=materialised, truncated=truncated, elapsed_ms=elapsed_ms, ) # ---- internals ---------------------------------------------------- @staticmethod def _clean(sql: str) -> str: sql = (sql or "").strip() # Strip surrounding markdown fence in case the LLM ignored the prompt. if sql.startswith("```"): sql = sql.strip("`") # After stripping backticks, drop a leading ``sql`` language tag. sql = re.sub(r"^\s*sql\b", "", sql, flags=re.IGNORECASE).strip() # Drop single trailing semicolon — we'll add LIMIT before it otherwise. sql = sql.rstrip(";").strip() return sql @staticmethod def _validate(sql: str) -> None: if not sql: raise SqlToolError("Empty SQL produced.") if not _SELECT_PATTERN.match(sql): raise SqlToolError( "Only SELECT / WITH statements are allowed. " f"Got: {sql.split()[0] if sql else '?'}\u2026" ) if _MULTI_STATEMENT.search(sql): raise SqlToolError("Multiple statements are not allowed.") if _FORBIDDEN_KEYWORDS.search(sql): raise SqlToolError( "SQL contains a forbidden keyword (write / DDL operation)." ) def _inject_limit(self, sql: str) -> str: if _LIMIT_PATTERN.search(sql): return sql return f"{sql}\nLIMIT {self.row_limit}" def _safe_url(self) -> str: # Hide credentials in log lines. url = self.database_url if "://" in url and "@" in url: scheme, rest = url.split("://", 1) creds, host = rest.rsplit("@", 1) if ":" in creds: user, _ = creds.split(":", 1) return f"{scheme}://{user}:****@{host}" return url def _describe_schema(self) -> str: try: inspector = inspect(self._engine) tables = sorted(inspector.get_table_names(schema="public")) except SQLAlchemyError as exc: raise SqlToolError(f"Cannot introspect schema: {exc}") from exc if not tables: return "(no tables found in 'public' schema)" chunks: list[str] = [] for table in tables: cols = inspector.get_columns(table, schema="public") pk = inspector.get_pk_constraint(table, schema="public").get( "constrained_columns", [] ) or [] fks = inspector.get_foreign_keys(table, schema="public") or [] col_lines: list[str] = [] for c in cols: pk_marker = " PRIMARY KEY" if c["name"] in pk else "" nullable = "" if c.get("nullable", True) else " NOT NULL" col_lines.append(f" {c['name']} {c['type']}{nullable}{pk_marker}") fk_lines: list[str] = [] for fk in fks: local = ", ".join(fk["constrained_columns"]) remote_table = fk["referred_table"] remote_cols = ", ".join(fk["referred_columns"]) fk_lines.append(f" FOREIGN KEY ({local}) -> {remote_table}({remote_cols})") block = f"TABLE {table} (\n" + ",\n".join(col_lines) if fk_lines: block += "\n --\n" + "\n".join(fk_lines) block += "\n)" chunks.append(block) return "\n\n".join(chunks)