""" 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()) )