src/agents/tools/sql_query.py

"""
Structured data query tool — CSV to SQLite with text-to-SQL.

Loads CSV files into an in-memory SQLite database and lets the agent
write SQL queries to answer questions about structured data.  This
covers the *structured knowledge* modality that enterprise agents
need alongside document retrieval (unstructured) and web search (live).

Safety: read-only mode, SELECT-only validation, row-count cap, timeout.
"""

from __future__ import annotations

import csv
import io
import sqlite3
from pathlib import Path

from src.agents.schemas import Tool, ToolParam

_MAX_ROWS = 50
_TIMEOUT_SECONDS = 5


class SQLStore:
    """In-memory SQLite database that ingests CSVs and executes read-only queries."""

    def __init__(self) -> None:
        self._conn = sqlite3.connect(":memory:", check_same_thread=False)
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._tables: dict[str, str] = {}  # table_name -> CREATE TABLE DDL

    @property
    def tables(self) -> dict[str, str]:
        """Map of table name to its CREATE TABLE statement."""
        return dict(self._tables)

    @property
    def num_tables(self) -> int:
        return len(self._tables)

    def ingest_csv(self, name: str, text: str) -> str:
        """Load CSV text into a table, return the schema description."""
        # Sanitise table name
        table_name = "".join(c if c.isalnum() or c == "_" else "_" for c in name)
        # Remove .csv suffix (as underscored form) if present
        for suffix in ("_csv", ".csv"):
            if table_name.endswith(suffix):
                table_name = table_name[: -len(suffix)]
                break
        table_name = table_name.strip("_") or "data"

        reader = csv.reader(io.StringIO(text))
        headers = next(reader)
        # Clean column names
        columns = [
            "".join(c if c.isalnum() or c == "_" else "_" for c in h.strip())
            for h in headers
        ]

        # Infer types from first row
        first_row = next(reader, None)
        col_types: list[str] = []
        for i, col in enumerate(columns):
            if first_row and i < len(first_row):
                val = first_row[i].strip()
                try:
                    int(val)
                    col_types.append("INTEGER")
                except ValueError:
                    try:
                        float(val)
                        col_types.append("REAL")
                    except ValueError:
                        col_types.append("TEXT")
            else:
                col_types.append("TEXT")

        # Create table
        col_defs = ", ".join(
            "{} {}".format(c, t) for c, t in zip(columns, col_types)
        )
        ddl = "CREATE TABLE {} ({})".format(table_name, col_defs)
        self._conn.execute("DROP TABLE IF EXISTS {}".format(table_name))
        self._conn.execute(ddl)
        self._tables[table_name] = ddl

        # Insert rows (including the first row we already read)
        placeholders = ", ".join("?" for _ in columns)
        insert_sql = "INSERT INTO {} VALUES ({})".format(table_name, placeholders)

        rows: list[tuple] = []
        if first_row:
            rows.append(tuple(first_row))
        for row in reader:
            if row:
                rows.append(tuple(row))
        if rows:
            self._conn.executemany(insert_sql, rows)
        self._conn.commit()

        return "Table '{}': {} rows, {} columns ({})".format(
            table_name, len(rows), len(columns), ", ".join(columns),
        )

    def execute_query(self, sql: str) -> str:
        """Execute a read-only SQL query and return formatted results."""
        # Safety: only allow SELECT and PRAGMA (for schema discovery)
        stripped = sql.strip().upper()
        if not (stripped.startswith("SELECT") or stripped.startswith("PRAGMA")):
            return "Error: only SELECT and PRAGMA queries are allowed."

        try:
            cursor = self._conn.execute(sql)
            col_names = [desc[0] for desc in cursor.description]
            rows = cursor.fetchmany(_MAX_ROWS)

            if not rows:
                return "Query returned 0 rows."

            # Format as a readable table
            lines = [" | ".join(col_names)]
            lines.append("-" * len(lines[0]))
            for row in rows:
                lines.append(" | ".join(str(v) for v in row))

            result = "\n".join(lines)
            total = cursor.fetchone()
            if total is not None:
                result += "\n... (showing first {} rows)".format(_MAX_ROWS)
            return result

        except sqlite3.Error as e:
            return "SQL error: {}".format(e)

    @staticmethod
    def _column_types(ddl: str) -> list[tuple[str, str]]:
        """Parse (name, type) pairs from a CREATE TABLE DDL."""
        inner = ddl.split("(", 1)[1].rsplit(")", 1)[0] if "(" in ddl else ""
        out: list[tuple[str, str]] = []
        for c in inner.split(","):
            parts = c.strip().split()
            if len(parts) >= 2:
                out.append((parts[0], parts[1]))
        return out

    def _distinct_values(self, table: str, col: str, limit: int = 40) -> list[str]:
        """Distinct values of a column (up to limit+1, to detect overflow)."""
        rows = self._conn.execute(
            "SELECT DISTINCT {} FROM {} LIMIT {}".format(col, table, limit + 1)
        ).fetchall()
        return [str(r[0]) for r in rows if r[0] is not None]

    def get_schema_description(self) -> str:
        """Describe all tables for the system prompt.

        Beyond column names, this lists the distinct values of low-cardinality
        TEXT columns and an example query.  Small models otherwise treat a
        long/tidy table's category values (e.g. a ``metric`` of ``total_revenue``)
        as column names and produce "no such column" errors.
        """
        if not self._tables:
            return "No data tables loaded."
        parts: list[str] = []
        for table_name, ddl in self._tables.items():
            coltypes = self._column_types(ddl)
            row_count = self._conn.execute(
                "SELECT COUNT(*) FROM {}".format(table_name)
            ).fetchone()[0]
            col_str = ", ".join("{} {}".format(c, t) for c, t in coltypes)
            parts.append("{} ({} rows): {}".format(table_name, row_count, col_str))

            numeric_cols = [c for c, t in coltypes if t in ("INTEGER", "REAL")]
            example_filters: list[tuple[str, str]] = []
            cat_lines: list[str] = []
            for c, t in coltypes:
                if t != "TEXT":
                    continue
                vals = self._distinct_values(table_name, c, limit=40)
                if 1 <= len(vals) <= 40:
                    cat_lines.append("    {}: {}".format(c, ", ".join(vals)))
                    if len(example_filters) < 2:
                        example_filters.append((c, vals[0]))
            if cat_lines:
                parts.append("  filter these columns by exact value:")
                parts.extend(cat_lines)
            if numeric_cols and example_filters:
                where = " AND ".join(
                    "{}='{}'".format(c, v) for c, v in example_filters
                )
                parts.append(
                    "  example: SELECT {} FROM {} WHERE {}".format(
                        numeric_cols[0], table_name, where,
                    )
                )
        return "\n".join(parts)

    def close(self) -> None:
        self._conn.close()


def make_sql_tool(store: SQLStore) -> Tool:
    """Create a SQL query tool bound to the given SQLStore."""

    def _get_description() -> str:
        schema = store.get_schema_description()
        return (
            "Execute a SQL SELECT query against uploaded CSV data. "
            "Always run a query to get the answer — do not guess from the schema. "
            "Tables may be in long format: filter the listed columns by their exact "
            "values (see 'filter these columns by exact value') and SELECT the numeric "
            "column — do not use a category value as a column name.\n"
            "Tables:\n{}".format(schema)
        )

    def execute(query: str) -> str:
        return store.execute_query(query)

    tool = Tool(
        name="query_data",
        description=_get_description(),
        params=[
            ToolParam(
                "query", "str",
                "A SQL SELECT query to execute, e.g. 'SELECT AVG(salary) FROM employees'.",
            ),
        ],
        execute=execute,
    )
    # Attach store reference so the description can be refreshed after new CSV ingestion
    tool._sql_store = store  # type: ignore[attr-defined]
    return tool


def refresh_sql_tool_description(tool: Tool) -> None:
    """Update the tool description after new CSVs are ingested."""
    store: SQLStore = tool._sql_store  # type: ignore[attr-defined]
    tool.description = (
        "Execute a SQL SELECT query against uploaded CSV data. "
        "Always run a query to get the answer — do not guess from the schema. "
        "Tables may be in long format: filter the listed columns by their exact "
        "values (see 'filter these columns by exact value') and SELECT the numeric "
        "column — do not use a category value as a column name.\n"
        "Tables:\n{}".format(store.get_schema_description())
    )