Skip to content

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Latest commit

 

History

4 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 

Repository files navigation

NL2SQL Agent for BigQuery

This project covers two versions of the same idea: ask a question in English, get BigQuery SQL, get an answer. v1 is a fixed pipeline with the table schema pasted into the prompt. v2 hands Gemini two tools and lets it decide what to call and in what order. I asked both the identical question, "What is the average fare and average trip distance across every taxi trip ever recorded?", against the public chicago_taxi_trips dataset.

The result worth reporting is not that they disagreed. They agreed. v1 returned an average fare of 13.816667710136851 and an average distance of 3.486804397836751, after silently restricting the data to trips between 2013-01-01 and 2024-01-01, a range it invented. v2 pulled the real schema first, wrote no date filter at all, and answered $13.82 and 3.49 miles. The same numbers to the cent and the hundredth of a mile.

That is what makes v1's behavior worth writing about. Nothing in its output would have told me a filter was there. I only caught it because I read the SQL.

The cause is one line of my own prompt. v1's instructions say Always include a date filter on trip_start_timestamp to limit scanned data. The model had no way to know the table's real date range, so it satisfied the rule the only way available to it: it made one up. The instruction was followed exactly. The answer was still built on a premise the model fabricated and never disclosed.

I built this on my own time to get up to speed on the Gemini API and BigQuery. Solo, my own idea, public dataset, not coursework.

The two versions

v1 (src/agent_v1.py) v2 (src/agent_v2.py)
Shape Fixed pipeline, one generation call Chat session with tool calling
Schema Hardcoded as a string in the prompt Fetched live via bq.get_table()
Control flow Python decides: generate, estimate, execute Model decides which tool to call and when
BigQuery access Python calls it directly Only through the run_query tool
Date filter on the all-time question 2013-01-01 to 2024-01-01, invented None, correctly
Bytes scanned on that question 5.11 GB not separately measured

v2 exposes exactly two functions to the model:

def get_table_schema() -> str:
    """Returns the column names, types, and row count for the Chicago taxi trips table."""

def run_query(sql: str) -> str:
    """Runs a BigQuery SQL query and returns the rows.
    Refuses any query that would scan more than 10 GB."""

Gemini reads those docstrings and signatures as the tool definitions. The system instruction tells it to call get_table_schema first and never to guess date ranges.

Prompt rules versus code enforcement

The whole project is really one thesis with a failure case and a fix.

Both versions carry the same kind of constraint, a rule about what the agent is allowed to do. They enforce it in two different places, and only one of those placements actually holds.

The failure case is a rule in the prompt. v1 asks for a date filter in English. The model produced one, and because it had no real date range to work from, the value was invented. There was no step in the system capable of noticing. The prompt cannot check its own output, and the fabricated range flowed straight into a query I then ran and believed. A rule stated in a prompt is a request. A request can be satisfied in ways that defeat the purpose of asking.

The fix is the same kind of rule written in Python. The cost ceiling never appears in any prompt. It is a dry run and a comparison:

cfg = bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
gb = bq.query(sql, job_config=cfg).total_bytes_processed / 1e9

if gb > MAX_GB:
    return f"REFUSED: would scan {gb:.2f} GB, over the {MAX_GB} GB limit."

The model cannot argue past this one, and the reason is structural rather than behavioral. In v2 the model never touches BigQuery. run_query is the only path from the model to the warehouse, and the check sits inside that function, ahead of execution. There is no phrasing of a request that reaches the billing meter without passing through this if. The model can be as confident, insistent, or confused as it likes and the ceiling is unchanged.

What happened when it fired. I went looking for a query wide enough to break the ceiling and asked for one: "For every trip in the dataset, show me the pickup and dropoff location, the fare, the tip, the company, and the taxi id. Just the first 100."

Gemini wrote the obvious query, six columns with LIMIT 100, and run_query dry ran it at 49.37 GB. The guard refused it. Three things followed, in order.

  1. The model tried to shrink the scan with SAMPLE SYSTEM, which is not valid BigQuery. The dry run raised, and my except block handed the syntax error back to it as a string.
  2. It read that string and corrected itself to TABLESAMPLE SYSTEM. I never wrote retry logic. The error text was already sitting in the conversation as a tool result, so the model treated it as one.
  3. TABLESAMPLE SYSTEM (1 PERCENT) with LIMIT 100 dry ran at 0.46 GB and executed.

Then it told me what it had done: "sampled using table sampling to stay well within the query limits."

That last line is the part I did not expect, and it is why this run is in the README at all. Enforcing the rule in code produced disclosure as a side effect. v1's invented date range stayed hidden because nothing in the system ever pushed back on it, and a prompt rule can be satisfied quietly. The refusal was different in kind. It was an external event the model had to react to, and reacting to it meant accounting for a change of plan that something outside itself had forced. My system instruction asks the agent to show the SQL it ran. It never asks the agent to explain why the SQL changed shape.

I would not oversell this. One run, one model, and the disclosure is a sentence I could easily have skimmed past. But it sharpens the thesis rather than replacing it: a constraint in code stops the thing you cannot afford, and it also tends to leave a trace in the answer, because the model has to narrate its way around an obstacle it cannot remove.

The uncomfortable part. v2's fix for the hallucination is itself a prompt rule. Always call get_table_schema first and Never guess date ranges live in the system instruction, which puts them in exactly the category I just described as unreliable. v2 got the right answer on this question, once. I have no guarantee it always will.

So the honest reading of my own project is narrower than the demo makes it look. Giving the model a schema tool did not make it obey. It removed the vacuum that made fabrication the path of least resistance, which is a real improvement and a different thing from a guarantee. The only rule here I would actually defend as enforced is the 10 GB check, because it is the only one the model cannot reach around.

The general form: give the model tools to make good behavior easy, and put the constraints you genuinely cannot afford to have violated in code, on the narrow path between the model and the thing that costs money or does damage.

What BigQuery charged me

BigQuery bills on bytes scanned, not rows returned, and the gap between those is larger than I expected.

Query Result size Dry run estimate
Payment type counts for one month (src/test_bq.py) 7 rows 3.36 GB in the console, 3.61 GB from the Python client
v1's all-time averages, with its invented date range 1 row 5.11 GB
Busiest companies 5 rows 5.76 GB
Six columns for the first 100 trips, no filter (v2, first attempt) 100 rows asked for 49.37 GB, refused
The same question after the model added TABLESAMPLE SYSTEM (1 PERCENT) 100 rows 0.46 GB

The last two rows are the clearest illustration of the gap. I asked for 100 rows and wrote LIMIT 100 in the SQL, and BigQuery priced it at 49.37 GB, because six unfiltered columns across the full table is what it has to read before it can throw almost all of it away. The payment type query shows the same thing on the WHERE side: 7 rows out, over 3 GB scanned. WHERE and LIMIT both bound what comes back, not what gets read. Sampling was the only one of the three that moved the number, from 49.37 GB to 0.46 GB, because it changes what gets read in the first place.

Two caveats I am leaving in rather than tidying away. The console and the Python client gave different estimates for the same query, 3.36 versus 3.61 GB, and I did not chase down which is authoritative or why they differ. And the mechanism I would offer for the WHERE behavior, that BigQuery bills for bytes read out of the columns a query references so a filter on a non-partitioned column still reads that column end to end, is something I have read rather than something I measured here. The behavior is measured. The explanation is not.

The 10 GB guard has now fired once, on that 49.37 GB query. Until then the largest query I had measured was 5.76 GB, so the threshold sat above everything I ran by accident rather than by design. What the refusal path did when it finally executed is written up in the section above.

Running it

Requires Python 3, a GCP project with the BigQuery API enabled, and a Gemini API key.

pip install google-genai google-cloud-bigquery

gcloud auth application-default login
export GOOGLE_CLOUD_PROJECT=your-project-id
export GEMINI_API_KEY=your-key

python src/test_bq.py       # BigQuery reachable, and the cost probe
python src/test_gemini.py   # Gemini reachable
python src/agent_v1.py      # the pipeline
python src/agent_v2.py      # the tool-calling agent

bigquery.Client() and genai.Client() both pick up credentials from the environment, so no keys appear anywhere in this repository.

The question is a variable at the bottom of each agent file. Edit it there to ask something else.

requirements.txt is a full freeze of the virtualenv I used, so the versions are the exact ones these runs were made with. The only direct dependencies are google-genai and google-cloud-bigquery.

What broke

Repeated 503s. gemini-3.7-flash returned 503 over and over while I was getting the first connectivity script working. I moved to gemini-3.5-flash-lite and it has been stable since. src/test_gemini.py still names gemini-3.7-flash, which I left alone because it is where the problem showed up.

Stale model names. The names I reached for from memory were not current. I stopped guessing and read Google's model documentation, which is also how I picked the replacement. Small lesson, but it is the one that cost me the most time relative to how interesting it was.

Limitations

  • One question, run once. The v1 versus v2 comparison rests on a single question and a single execution of each. It is an observation, not a benchmark, and I do not know how often v1 fabricates a range or how often v2 avoids it.
  • The agreement is not proof of correctness. v1 and v2 matching to two decimal places tells me the invented 2013 to 2024 window happens to cover nearly all the data. It does not tell me either answer is right, because I never verified them against an independent count.
  • No evaluation harness. Correctness was checked by reading the generated SQL myself. There are no assertions and no known-good answers anywhere in this repo.
  • test_bq.py and test_gemini.py are not tests. They are connectivity probes with no assertions, kept under their original names because that is what I called them while working.
  • It decided on a 1 percent sample without asking me. When the guard refused the 49.37 GB query, the model picked table sampling on its own, ran it, and told me afterward. It disclosed the substitution, which is more than v1 ever did for its invented date range, but disclosure after the fact is not consent. Nothing in the design gives the agent a way to stop and ask whether a sampled answer is an acceptable answer to the question I asked. On any question where sampling changes what the result means, the only thing between me and a quietly different answer is whether I read the disclosure line.
  • Single table, single model. Everything here is one public table and one Gemini model. Nothing about multi-table joins, ambiguous column names, or model comparison was explored.
  • No retry logic, and the recovery I got was incidental. Handing an invalid SQL error string back to the model is the whole of v2's error handling. On the SAMPLE SYSTEM mistake it happened to be enough, and the model fixed its own syntax. That is a lucky outcome, not a designed one. There is no attempt counter and no cap, so nothing in this code would stop a model that kept producing the same bad query from looping on it.

What I would do next

Make the assumption visible in the answer. The cheapest fix for the entire failure described above is to require the agent to report the date range it used as part of every answer. A fabricated filter that has to be printed next to the number is a fabrication I catch in one second instead of by reading SQL.

Build a question set with known answers. Ten to twenty questions with independently verified results, run against both versions, so I can put a rate on the hallucination instead of an anecdote.

Log every query and its estimate. One CSV row per generated query, with the SQL and the dry run figure, so cost is measured across a session rather than remembered from the ones I happened to notice.

Make the agent ask before it substitutes. The 1 percent sample should have been a question put to me, not a decision made for me. A refusal from run_query could come back with an instruction to stop and propose the workaround rather than run it, which puts the choice where the cost and the meaning of the answer both land. The same change wants an attempt cap, so a model that cannot find a query under the ceiling gives up instead of grinding.

Understand the scan cost properly. Check whether the table is partitioned, measure what selecting fewer columns does to the estimate, and replace the explanation I borrowed above with one I verified.

Repository contents

src/agent_v1.py      fixed pipeline: hardcoded schema, one generation, dry run, execute
src/agent_v2.py      tool-calling agent: schema tool and query tool, model picks the order
src/test_bq.py       first BigQuery connectivity check, plus the dry run cost probe
src/test_gemini.py   first Gemini connectivity check, still naming the model that 503'd
requirements.txt     full freeze of the environment these runs were made in

The code is published as it ran, including the unused import os in agent_v1.py and the probe filenames.

License and data

MIT, see LICENSE.

The data is bigquery-public-data.chicago_taxi_trips.taxi_trips, part of the Google Cloud public datasets program and sourced from the City of Chicago's open data portal. No data is redistributed here. Running any of these scripts queries that table under your own GCP billing account, and the dry run estimates above are what those queries cost in bytes scanned.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages