Text-to-SQL Agent: Answer Questions From a Database Without Letting the Model Touch the Data

"""Ask Tidewell Gym's database questions in plain English: the model writes SQL, the code checks and runs it."""
import json
import os
import re
import sqlite3
import time

import build_db
import harness

DB = "gym.db"
MAX_ROWS = 50
TIME_LIMIT = 2.0   # seconds per query


def ensure_db():
    if not os.path.exists(DB):
        build_db.build(DB)


# ---------- Step 1: describe the database ----------
def schema_text(conn):
    """For every table (in sqlite_master order): its CREATE statement exactly as stored (with its comments), its
    row count, and its first 2 rows as tuples."""
    # TODO (Step 1): loop over conn.execute("SELECT name, sql FROM sqlite_master WHERE type = 'table' ORDER BY rowid");
    # for each table give its CREATE sql, its COUNT(*) and its first 2 rows (LIMIT 2). Join the tables with blank lines.
    raise NotImplementedError("Step 1: write schema_text()")


# ---------- Step 2: a connection that can only read ----------
ALLOWED_ACTIONS = {sqlite3.SQLITE_SELECT, sqlite3.SQLITE_READ, sqlite3.SQLITE_FUNCTION, sqlite3.SQLITE_RECURSIVE}


def authorize(action, arg1, arg2, db_name, trigger):
    """sqlite3 calls this for every action a statement needs while it is compiled."""
    return sqlite3.SQLITE_OK


def connect_readonly():
    """A connection that opens gym.db read-only and denies every action outside ALLOWED_ACTIONS."""
    return sqlite3.connect(DB, check_same_thread=False)


def run_query(sql, max_rows=MAX_ROWS):
    """Run one query on a read-only connection. Returns {"columns", "rows" (lists), "truncated"} or {"error": text}.
    Stop queries that run longer than TIME_LIMIT seconds; return at most max_rows rows."""
    conn = sqlite3.connect(DB)
    try:
        cur = conn.execute(sql)
        rows = cur.fetchall()
        conn.commit()
        return {"columns": [d[0] for d in cur.description or []], "rows": [list(r) for r in rows], "truncated": False}
    except sqlite3.Error as e:
        return {"error": str(e)}
    finally:
        conn.close()


# ---------- Step 3: check, run, repair ----------
def extract_sql(reply):
    """The SQL inside the reply's ```sql block (or the whole reply if there is none), stripped."""
    m = re.search(r"```(?:sql)?\s*(.*?)```", reply or "", re.S | re.I)
    return (m.group(1) if m else reply or "").strip()


def check_sql(sql):
    """None if the SQL is one read-only statement that compiles; else the reason it is not."""
    return None


def build_prompt(question, schema):
    prompt = ("You write SQLite queries for Tidewell Gym's database. Read the comments in the schema: they hold "
              "the business rules.\n\n" + schema)
    return prompt + ("\n\nAnswer with one SQLite SELECT query in a ```sql block and nothing else. If the request "
                     "asks to change data, still answer with a SELECT query that reads what is relevant.")


def ask(question, max_repairs=2):
    """Returns {"sql", "columns", "rows", "error", "attempts"}: the last SQL tried, its result (or error)."""
    conn = sqlite3.connect(f"file:{DB}?mode=ro", uri=True)
    schema = schema_text(conn)
    conn.close()
    messages = [{"role": "system", "content": build_prompt(question, schema)}, {"role": "user", "content": question}]
    for attempt in range(1, max_repairs + 2):
        reply = harness.complete(messages, max_tokens=500)["content"] or ""
        sql = extract_sql(reply)
        problem = check_sql(sql)
        result = run_query(sql) if problem is None else {"error": problem}
        return {"sql": sql, "columns": result.get("columns", []), "rows": result.get("rows", []),
                "error": result.get("error"), "attempts": attempt}


# ---------- Step 4: grade by results ----------
def norm(v):
    """Numbers rounded to 2 decimal places (as floats), everything else unchanged."""
    return round(float(v), 2) if isinstance(v, (int, float)) and not isinstance(v, bool) else v


def same_result(expected, got):
    """True if got answers the question as well as expected: the same number of rows, and every expected row matches
    a different got row that contains all its values (compared with norm()). Row order, column order, column names
    and extra columns in got do not matter."""
    return expected == got


# ---------- Step 5: examples that teach the rules ----------
EXAMPLES = json.load(open("examples.json", encoding="utf-8"))


def words(text):
    return set(re.findall(r"[a-z0-9]+", text.lower())) - {"the", "a", "an", "in", "of", "how", "what", "is", "are",
                                                          "was", "were", "many", "much", "which", "do", "does", "did"}


def pick_examples(question, k=4):
    """The k examples whose questions share the most words with this question (Jaccard similarity of words()),
    most similar first; ties keep the library order."""
    return []