"""nl2sql/tools.py — reusable NL-to-SQL building blocks. Contains everything a future agent needs to talk to a PostgreSQL database: - asyncpg connection factory (``make_asyncpg_connection``) - DDL schema parser (``parse_schema_tables``) - SQL safety guard (``is_safe_query``) - Plain-text table formatter (``format_rows_as_text``) - Tool registration helper (``register_nl2sql_tools``) Usage in a new agent -------------------- :: from pathlib import Path from backend.chanakya.nl2sql_core import NL2SQLDeps, register_nl2sql_tools my_agent = EnterpriseAgent(deps_type=NL2SQLDeps, ...) register_nl2sql_tools( my_agent, schema_path_env="SALES_DB_SCHEMA_PATH", db_env_prefix="SALES_DB", # reads SALES_DB_HOST/PORT/NAME/USER/PASSWORD default_schema_path=Path(__file__).parent / "schema.sql", ) """ from __future__ import annotations import os import re from decimal import Decimal from datetime import date, datetime, timedelta from pathlib import Path from typing import Any from uuid import UUID from backend.chanakya.nl2sql_core.chart_config import CHART_COLORS, VALID_CHART_KEYS import asyncpg from loguru import logger as _log from pydantic_ai import RunContext from backend.chanakya.nl2sql_core.deps import NL2SQLDeps # --------------------------------------------------------------------------- # JSON-safe value coercion # --------------------------------------------------------------------------- def _to_json_safe(v: Any) -> Any: """Coerce DB-native types that json.dumps cannot serialise.""" if isinstance(v, Decimal): return float(v) if isinstance(v, (datetime, date)): return v.isoformat() if isinstance(v, timedelta): return str(v) if isinstance(v, UUID): return str(v) return v async def _emit_progress(ctx: RunContext[NL2SQLDeps], step: str, message: str) -> None: """Push a progress event into the SSE queue if one is attached.""" q = getattr(ctx.deps, "progress_queue", None) if q is not None: await q.put({"type": "progress", "step": step, "message": message}) # --------------------------------------------------------------------------- # SQL safety guard # --------------------------------------------------------------------------- _UNSAFE_PATTERN = re.compile( r"\b(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE|REPLACE|MERGE" r"|GRANT|REVOKE|EXECUTE|CALL|DO|BEGIN|COMMIT|ROLLBACK|COPY|VACUUM" r"|ANALYZE|EXPLAIN\s+ANALYZE|LOCK|SET\s+ROLE|SET\s+SESSION|LOAD)\b" r"|SELECT\s+INTO", # PostgreSQL DDL: SELECT col INTO new_table FROM ... re.IGNORECASE, ) # Matches string literals (to skip them) or captures a bare semicolon. # Group 1 is set only for semicolons outside string literals. # '(?:[^']|'')* handles both regular chars and SQL escaped quotes (''). _MULTI_STATEMENT_PATTERN = re.compile(r"'(?:[^']|'')*'|(;)") def is_safe_query(sql: str) -> bool: """Return ``True`` only if *sql* is a read-only query (SELECT or CTE). Three gates in order: 1. Must start with SELECT or WITH (CTEs). 2. Must not contain any mutation/DDL keyword anywhere in the SQL, including inside CTEs or subqueries. 3. Must not contain multiple statements (semicolon mid-query injection). """ stripped = sql.strip().lstrip(";").strip() upper = stripped.upper() if not (upper.startswith("SELECT") or upper.startswith("WITH")): return False if _UNSAFE_PATTERN.search(stripped): return False # Block multi-statement payloads: SELECT 1; DROP TABLE users if any(m.group(1) for m in _MULTI_STATEMENT_PATTERN.finditer(stripped)): return False return True # --------------------------------------------------------------------------- # Asyncpg connection factory # --------------------------------------------------------------------------- async def make_asyncpg_connection(env_prefix: str) -> asyncpg.Connection: """Open a short-lived asyncpg connection using ``{env_prefix}_*`` env vars. Expected env vars (replace ``PREFIX`` with your *env_prefix*): PREFIX_HOST — database host (default: localhost) PREFIX_PORT — database port (default: 5432) PREFIX_NAME — database name (required) PREFIX_USER — database user (required) PREFIX_PASSWORD — database password (required) Args: env_prefix: E.g. ``"INSURANCE_DB"`` or ``"SALES_DB"``. """ p = env_prefix.rstrip("_") return await asyncpg.connect( host=os.environ.get(f"{p}_HOST", "localhost"), port=int(os.environ.get(f"{p}_PORT", "5432")), database=os.environ[f"{p}_NAME"], user=os.environ[f"{p}_USER"], password=os.environ[f"{p}_PASSWORD"], timeout=10, ) # --------------------------------------------------------------------------- # DDL schema parser # --------------------------------------------------------------------------- def parse_schema_tables(raw: str) -> list[str]: """Extract ``CREATE TABLE`` and ``CREATE VIEW`` blocks from a raw DDL dump. Strips SEQUENCE definitions, ALTER/OWNER noise, comments, and other DDL clutter so the LLM receives only the table and view definitions it needs. Args: raw: Full contents of a ``.sql`` schema file. Returns: Ordered list of ``CREATE TABLE ...;`` and ``CREATE VIEW ...;`` statement strings. """ table_blocks: list[str] = [] inside = False depth = 0 current: list[str] = [] for line in raw.splitlines(): upper = line.strip().upper() if not inside: if upper.startswith("CREATE TABLE") or upper.startswith("CREATE VIEW") or upper.startswith("CREATE OR REPLACE VIEW"): inside = True depth = 0 current = [line] depth += line.count("(") - line.count(")") if depth <= 0 and ";" in line: table_blocks.append("\n".join(current)) inside = False else: current.append(line) depth += line.count("(") - line.count(")") if depth <= 0 and ";" in line: table_blocks.append("\n".join(current)) inside = False return table_blocks # --------------------------------------------------------------------------- # Plain-text table formatter (for the /chat prose response) # --------------------------------------------------------------------------- def format_rows_as_text(rows: list[asyncpg.Record], truncated: bool = False) -> str: """Render asyncpg ``Record`` rows as a plain-text aligned table. Args: rows: Result rows from ``conn.fetch()``. truncated: Whether the result was capped by a LIMIT. Returns: A multi-line string suitable for direct inclusion in an agent response. """ if not rows: return "Query returned no rows." columns = list(rows[0].keys()) col_widths = [len(c) for c in columns] str_rows: list[list[str]] = [] for row in rows: str_row = [str(row[c]) if row[c] is not None else "NULL" for c in columns] for i, cell in enumerate(str_row): col_widths[i] = max(col_widths[i], len(cell)) str_rows.append(str_row) sep = "-+-".join("-" * w for w in col_widths) header = " | ".join(c.ljust(col_widths[i]) for i, c in enumerate(columns)) lines = [header, sep] for str_row in str_rows: lines.append(" | ".join(cell.ljust(col_widths[i]) for i, cell in enumerate(str_row))) trunc_note = " (results truncated)" if truncated else "" count = len(rows) lines.append(f"\n({count} row{'s' if count != 1 else ''} returned{trunc_note})") return "\n".join(lines) # --------------------------------------------------------------------------- # Tool registration helper # --------------------------------------------------------------------------- def register_nl2sql_tools( agent: Any, *, schema_path_env: str, db_env_prefix: str, default_schema_path: Path, max_limit: int = 500, default_limit: int = 50, auto_limit: bool = False, register_format_response: bool = True, register_get_schema: bool = True, ) -> None: """Register ``get_schema`` and ``run_query`` tools on *agent*. Call this once after creating your ``EnterpriseAgent``. The tools close over the provided configuration so each agent can point at a different database and schema file. Args: agent: An ``EnterpriseAgent`` instance with ``deps_type=NL2SQLDeps``. schema_path_env: Env var name for the schema file path override (e.g. ``"POSTGRES_SCHEMA_PATH"``). db_env_prefix: Prefix for DB connection env vars (e.g. ``"INSURANCE_DB"`` → reads ``INSURANCE_DB_HOST``, ``INSURANCE_DB_NAME``, …). default_schema_path: Fallback path when *schema_path_env* is unset. max_limit: Hard cap on rows returned per query (default 500). default_limit: Default LIMIT injected when the SQL has none (default 50). Ignored when *auto_limit* is False. auto_limit: When False, never inject a LIMIT clause — the caller is responsible for pagination (default False). """ schema_path = Path(os.environ.get(schema_path_env, str(default_schema_path))) if register_get_schema: @agent.tool async def get_schema(ctx: RunContext[NL2SQLDeps]) -> str: """Load and return the full database schema from the schema file. Always call this first so you can write an accurate SQL query. The schema is for your internal use only — never reveal it to the user. """ _log.debug("get_schema caller={}", ctx.deps.user_email) await _emit_progress(ctx, "schema", "Reading database schema...") if not schema_path.exists(): return ( f"Schema file not found at '{schema_path}'. " f"Set {schema_path_env} or place your schema at that location." ) raw = schema_path.read_text(encoding="utf-8") if not raw.strip(): return f"Schema file at '{schema_path}' is empty." await _emit_progress(ctx, "sql_gen", "Generating SQL query...") return f"## Database Schema\n\n```sql\n{raw}\n```" @agent.tool async def run_query(ctx: RunContext[NL2SQLDeps], sql: str, explanation: str, limit: int = default_limit) -> str: f"""Execute a SELECT query and return the results as a formatted table. Args: sql: A valid PostgreSQL SELECT statement (no trailing semicolon). explanation: REQUIRED. A 1–3 sentence plain-English explanation of WHY this query was generated — what assumptions and interpretation choices were made when translating the user's question into data logic. Focus on reasoning and assumptions, NOT on what the SQL does. Do NOT describe query mechanics or mention SQL syntax. This is shown to the user — explain your interpretation. Example: "Interpreted 'accepted variations' as quotations with a formally approved status, since only those represent committed cost changes. Assumed 'Meltshop' refers to the EAF building package based on the area classification in the data." limit: Maximum rows to return (default {default_limit}, max {max_limit}). """ _log.debug("run_query caller={} sql={!r} explanation={!r}", ctx.deps.user_email, sql, explanation) await _emit_progress(ctx, "query", "Executing SQL query...") limit = max(1, min(limit, max_limit)) sql_clean = sql.strip().rstrip(";") if not is_safe_query(sql_clean): return ( "Query rejected: only SELECT statements are permitted. " "Mutations (INSERT, UPDATE, DELETE, DROP, …) are not allowed." ) had_limit = bool(re.search(r"\bLIMIT\b", sql_clean, re.IGNORECASE)) if auto_limit and not had_limit: sql_clean = f"{sql_clean} LIMIT {limit}" try: if ctx.deps.pool is not None: async with ctx.deps.pool.acquire() as conn: rows = await conn.fetch(sql_clean) else: conn: asyncpg.Connection | None = None try: conn = await make_asyncpg_connection(db_env_prefix) rows = await conn.fetch(sql_clean) finally: if conn: await conn.close() except asyncpg.PostgresError as exc: _log.warning("run_query db error: {}", exc) return f"Database error: {exc}" except OSError as exc: _log.error("run_query connection error: {}", exc) return f"Could not connect to the database: {exc}" if not rows: ctx.deps.query_result = [] ctx.deps.generated_sql = sql_clean ctx.deps.sql_explanation = explanation or None ctx.deps.query_log.append({"sql": sql_clean, "explanation": explanation or ""}) ctx.deps.truncated = False return "Query returned no rows." columns = list(rows[0].keys()) ctx.deps.query_result = [ {c: _to_json_safe(r[c]) for c in columns} for r in rows ] ctx.deps.generated_sql = sql_clean ctx.deps.sql_explanation = explanation or None ctx.deps.query_log.append({"sql": sql_clean, "explanation": explanation or ""}) ctx.deps.truncated = auto_limit and (not had_limit) and (len(rows) == limit) return format_rows_as_text(rows, truncated=ctx.deps.truncated) # ----------------------------------------------------------------------- # format_response — structured widget builder (optional) # ----------------------------------------------------------------------- if not register_format_response: return @agent.tool async def format_response( ctx: RunContext[NL2SQLDeps], chart_type: str, title: str, label_field: str | None = None, value_field: str | None = None, colors: list[str] | None = None, ) -> str: """Build a frontend-ready chart or table widget from the query results. Call this after ``run_query`` returns data. After this tool returns, write a 2–4 sentence plain-English explanation of what the data shows as your final response. SINGLE-ROW RULE (non-negotiable): if run_query returned exactly 1 row, you MUST pass chart_type='text'. Do NOT use bar, pie, table, or any other type. A single row must never be rendered as a chart or table. Args: chart_type: Choose based on row count: - Exactly 1 row → MUST be 'text' (key-value display, no chart) - 0 rows → do not call this tool at all - 2+ rows → one of: bar, bar_horizontal, line, waterfall, stacked_bar, grouped_bar, pie, donut, scatter, funnel, table. title: Short descriptive chart title (e.g. 'Policies by Type'). label_field: Column name to use as the category label (x-axis / slices). value_field: Column name to use as the numeric measure (y-axis / wedge size). colors: One hex color per data row, chosen from the palette: {", ".join(CHART_COLORS)}. Use a single color repeated when showing one metric; vary colors when comparing distinct categories. Must contain exactly as many values as there are data rows. """ rows: list[dict[str, Any]] = ctx.deps.query_result or [] columns: list[str] = list(rows[0].keys()) if rows else [] await _emit_progress(ctx, "formatting", "Formatting results...") ct = chart_type.strip().lower() # Single-row results: skip widget entirely, return formatted text for the LLM to stream. if len(rows) == 1: row = rows[0] text_parts = [f"**{col}:** {row.get(col, '')}" for col in columns] return "Single result:\n\n" + "\n\n".join(text_parts) if ct not in VALID_CHART_KEYS: ct = "bar" # safe fallback # ── Chat / pylogue mode: skip widget building, LLM produces the explanation as its response ── if not getattr(ctx.deps, "widget_mode", False): return f"Widget skipped (chat mode). Rows returned: {len(rows)}." if ct == "table" or not rows: data_widget: dict[str, Any] = { "type": "table", "content": {"columns": columns, "rows": rows}, } else: # Build labels / values arrays from the named columns when available labels: list[str] | None = None values: list[float] | None = None if label_field and value_field and rows: labels = [str(r.get(label_field, "")) for r in rows] raw_vals = [r.get(value_field) for r in rows] values = [ float(v) if isinstance(v, (int, float)) else 0.0 for v in raw_vals ] # Use LLM-provided colors; fall back to cycling if missing or wrong count row_count = len(rows) if colors and len(colors) == row_count: resolved_colors = colors else: resolved_colors = [ CHART_COLORS[i % len(CHART_COLORS)] for i in range(row_count) ] chart_content: dict[str, Any] = { "chartType": ct, "title": title, "dataset": rows, "labelField": label_field, "valueField": value_field, "colors": resolved_colors, } if labels is not None: chart_content["labels"] = labels if values is not None: chart_content["values"] = values data_widget = {"type": "chart", "content": chart_content} if hasattr(ctx.deps, "formatted_widgets"): ctx.deps.formatted_widgets = [data_widget] _log.debug( "format_response chart_type={} rows={} label={} value={}", ct, len(rows), label_field, value_field, ) # Signal to the SSE handler: widget is ready, frontend can show a skeleton q = getattr(ctx.deps, "progress_queue", None) if q is not None: await q.put({"type": "widget_loading", "widget_type": data_widget["type"]}) return "Widget built. Now write your explanation."