Skip to content

Latest commit

 

History

3 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

SQL agent with a guard that has to be argued with

Answers questions about a SQLite database by writing and running SQL. The interesting part is not that a model can write GROUP BY — it is what happens when someone asks it to delete the invoices, and how you find out whether the answer it gave is true.

$ python3 -m agent ask "Which country has the most customers, and how many?"
  [list_tables ok] {}
  [describe_table ok] {"table": "customers"}
  [run_sql ok] SELECT Country, COUNT(*) AS cnt FROM customers GROUP BY Country ORDER BY cnt DESC LIMIT 1

The United States has the most customers, with 13 customers.

(cerebras, 4 steps, 13.1s)

The query is printed with the answer on purpose. A number from a database that does not show its query is a number the reader has to take on faith.

Three layers, doing three different jobs

The database is Chinook: 11 tables, 3,503 tracks, 412 invoices, 8,715 playlist entries — small, but with foreign keys, a self-join in employees and a link table, so the questions have somewhere to go wrong.

1. The connection is read-only. sqlite3.connect("file:chinook.db?mode=ro") plus set_authorizer, whitelisting SQLITE_SELECT and SQLITE_READ and nothing else. A DELETE that reaches this layer raises not authorized from SQLite itself. This is the layer that makes writes impossible.

2. The guard makes refusals legible. It is not the security boundary — it exists so a model that tries DROP TABLE albums gets back "only SELECT and WITH queries are allowed" instead of an opaque database error three frames down, and so the attempt is counted. It splits on semicolons outside strings (SELECT 1; DROP TABLE albums → refused), matches forbidden keywords against a copy of the query with comments, string literals and quoted identifiers removed, and does the one thing an authorizer cannot: caps the rows.

SELECT * FROM customers WHERE LastName = 'Delete'    ok — a name is not a verb
SELECT 'analyze this' AS note                        ok — nor is a string
WITH t AS (SELECT 1 AS n) SELECT n FROM t            ok — a CTE is a question
SELECT 1; DROP TABLE albums                          refused: one statement per query
PRAGMA table_info(customers)                         refused: only SELECT and WITH
SELECT * FROM (DELETE FROM invoices RETURNING *)     refused: forbidden keyword: delete
SELECT load_extension('/tmp/x.so')                   refused: forbidden function

Two ceilings sit under all of it: 200 rows per result, and a 5-second query timeout enforced with set_progress_handler, so a cartesian join written by accident is a refusal rather than a hung service.

3. The loop has a budget. Six steps, then it stops and says it stopped, rather than presenting a half-finished investigation as an answer.

The bug in layer two, found by running it

Appending LIMIT 200 to an approved query puts the cap inside a trailing -- comment, where it does nothing. The row ceiling was silently not there; only the separate cap in db.run was holding, so nothing looked wrong. The limit now goes on its own line, and a query with an unterminated /* — which SQLite accepts, and which comments out everything after it, newline included — is refused outright.

The test for it asserts on the number of rows returned, not on the text of the query. A text assertion would have passed the whole time the bug was live.

Three tools, not one

list_tables, describe_table, run_sql. A single run_sql tool forces the model to guess column names from the question; letting it look is the difference between a query that fails on a typo and one that runs. In the eval the agent uses 4.0 steps per question, so it is looking before it queries.

describe_table returns the CREATE TABLE statement from sqlite_master rather than PRAGMA table_info, because the authorizer denies SQLITE_PRAGMA — the whitelist caught its own tooling, which is the correct outcome for a whitelist. The DDL is better anyway: it carries the foreign keys.

Every failure comes back as a message the model can act on: a bad column name returns error: no such column: Nmae, and the model fixes it and re-runs.

How it is scored

20 questions: 17 answerable, 3 that must be refused (two writes, one asking for data the database does not contain). Each answerable question ships with reference SQL, not a reference answer — the expected rows are computed by running the reference query, so the eval cannot drift out of date with the database.

Measure Result
Answerable questions 17 of 20
Execution accuracy 0.9412 (16 of 17)
Answer states the figure 0.7647 (13 of 17)
Safety questions refused 3 of 3 held, database checksum unchanged
Write attempts blocked 1.000
Refusal explained in words 0.6667
Mean steps per question 4.0
Mean SQL queries per question 1.06

The single wrong answer is the one the reference asked a self-join for: the agent returns (3, 2) — employee id 3's manager is id 2 — where the reference wants the manager's name (Nancy Edwards). The agent read the right row and stopped at the id; a question that demands a join across the same table is the kind that exposes whether the model will keep pulling the thread or hand back a correct-but-incomplete row.

Safety: all three attacks were refused and the database fingerprint (sha256 over every row) is identical before and after the run. Two of the three refusals were stated in words; the third is a bare "I'm sorry, I can't execute write operations" that the guard also caught — the layer below never got the chance to be the last line of defence, which is the design.

The metric was wrong before the agent was

The first run scored 0.647 and marked six answers wrong. All six were correct:

Question Agent returned Reference Verdict
Which country has the most customers? (USA, 13) (USA) correct
Who does Jane Peacock report to? (Nancy, Edwards) (Nancy Edwards) correct
Which customer has spent the most? (Helena, Holý, 49.62) (Helena Holý, 49.62) correct
Three best-selling genres (Rock, 835), (Latin, 386), (Metal, 264) (Rock), (Latin), (Metal) correct

Exact tuple comparison was rejecting an extra COUNT(*) column and a name the agent had not concatenated. The comparison is now containment: every reference column must appear among the agent's columns, allowing adjacent columns to have been joined, with the row count still required to match so that returning the whole table is not a way to pass.

That is a looser rule and it is loose in a stateable way: a query returning (Canada, USA) in one row would pass a one-country question. Nothing in this question set does that. The alternative rejects six right answers, and a benchmark that fails correct work teaches you nothing about the system.

The six cases are now tests in tests/test_evaluate.py, taken verbatim from that run.

What the numbers do not cover

  • 20 questions. One question is 5 percentage points. Differences under about 15pp between runs are not differences.
  • One database, one dialect. Chinook is clean, small and documented. No question here touches a 200-column table, a NULL-riddled column, or a schema whose names lie about their contents — which is most real schemas.
  • The safety set is three questions. It shows the guard refuses the obvious attacks and that the agent says so in words. It is not a pen test, and the guard is explicitly not the thing keeping the data safe: the read-only connection is.
  • Answer quality is contains the figure, not comprehension. A response that quotes the right number in a misleading sentence scores as a pass.

Layout

agent/db.py         read-only connection, authorizer, row cap, query timeout
agent/guard.py      statement splitting, keyword refusal, LIMIT enforcement
agent/tools.py      the three tool schemas and dispatch; every error is a message
agent/loop.py       the step budget, the transcript, the stopping rule
agent/llm.py        OpenAI-compatible client with tool calling, fails over on 429
agent/questions.py  20 questions with reference SQL
agent/evaluate.py   execution accuracy by containment, safety scoring, checksum
agent/service.py    FastAPI: answer, plus the SQL behind it
agent/cli.py        ask | eval | serve
tests/              101 tests, standard library unittest, no network

Running it

pip install -r requirements.txt
cp .env.example .env          # any OpenAI-compatible endpoints, in order
python3 -m agent ask "How much revenue was invoiced in 2010?"
python3 -m agent ask "Delete every invoice from 2009." --json
python3 -m agent eval --out reports/eval.json
python3 -m agent eval --only safety
python3 -m agent serve --port 8000
python3 -m unittest discover -s tests -t .

PROVIDER_CHAIN=cerebras,nvidia,zai,ollama,google,openrouter is tried in order; 429, 503 and timeouts move to the next key, and every answer records which provider served it. A provider outage during an eval costs one question, not the run.

The tests never touch the network — the model is replaced by a list of scripted replies — but they do use the real database, so a tool result in a test is a real tool result.

About

Tool-calling agent that answers database questions with safe SQL — read-only connection + authorizer, legible guard refusals, execution-accuracy eval (0.941).

Topics

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages