GenAI Analytics Agent

Natural language to SQL is one of those demos that looks trivial in a notebook and breaks catastrophically in production. Getting a model to turn “what was our fill rate last month?” into a syntactically correct Redshift query is a solved problem. Getting it to do that reliably, safely, at scale, across a schema with hundreds of tables and PII-adjacent data — that’s a different engineering problem entirely.

This is a write-up of the production system I built at Amazon Pharmacy and the lessons it forced. It covers the full pipeline — from intent classification through schema retrieval, generation, validation, execution, and feedback capture — along with the privacy architecture that made it safe to deploy, and the reinforcement learning loop that made it get better over time.

Why Naive Approaches Break

The simplest version of NL-to-SQL is: stuff the entire schema into the system prompt and ask GPT-4 to write a query. For a schema with five tables and thirty columns, this works surprisingly well. For a production data warehouse with 300 tables and 8,000 columns — the kind of environment Uber’s QueryGPT operates in — it fails in four distinct ways.

Context window saturation. A full schema dump for a mature warehouse will exceed the context limit of any production model. Even within the limit, attention degrades at the periphery. Tables mentioned in the first third of the prompt will be used more reliably than those buried in the middle. You cannot solve a retrieval problem by throwing more tokens at it.

Column hallucination. When a model doesn’t find the column it needs, it invents one. This isn’t a failure mode you can catch at generation time — the SQL looks valid. It fails at execution, or worse, it executes on a column that happens to exist but means something different. In a pharmacy context, a hallucinated column mapping in a cost query isn’t just wrong; it could drive a procurement decision in the wrong direction.

Schema drift. Columns get renamed. Deprecated tables stay in the warehouse for months after they stop being populated. A prompt that was accurate in Q1 is subtly wrong by Q3 because nobody updated the schema description embedded in it. Static schema injection has no versioning story.

Ambiguity compounding. Business questions are semantically ambiguous in ways that domain experts don’t notice. “Active customers” means different things to finance, marketing, and operations. An LLM will make a silent choice. The query runs. The result looks plausible. The decision it informs is based on a definition nobody agreed on. Ambiguity in the input doesn’t produce an error — it produces a confidently wrong answer.

The Production Pipeline

A production NL-to-SQL system is not a prompt. It’s a pipeline with at least six distinct stages, each of which can fail independently and needs to be observable independently.

Intent classification is the first gate. Before schema retrieval, a lightweight classifier — either a small fine-tuned model or a zero-shot LLM call — determines whether the query is (a) answerable from the available data, (b) a meta-question about the system itself, or (c) out of scope. Uber’s QueryGPT uses a dedicated Intent Agent to route queries to the correct “workspace” — a curated collection of tables and SQL samples scoped to a specific business domain like Mobility, Ads, or Core Services. The workspace concept is the right abstraction: it bounds the retrieval problem and improves precision dramatically.

Schema retrieval happens against a pre-indexed embedding store. Tables and columns are embedded as rich semantic descriptions — not raw DDL. The query embedding is compared via cosine similarity to retrieve the most-relevant subset of schema, typically ten to twenty columns across three to five tables. This is where the difference between “user_id” and “unique identifier for a customer session, joins to the orders table via session_id” becomes load-bearing.

Validation runs before execution. A SQL parser checks syntax. A column-existence check runs the referenced column names against the live schema catalog. A dry-run against EXPLAIN (or Redshift’s EXPLAIN equivalent) catches join errors and type mismatches without touching actual data. Queries that fail validation are not executed — they are fed back to the generation stage with the error appended.

Execution runs in a read-only IAM role with no write permissions on any table. Row limits are enforced at the query wrapper layer, not the generated SQL — because a model will sometimes omit a LIMIT clause. Result formatting is a separate step: numbers get locale-appropriate formatting, column headers get humanized, and the result is accompanied by the SQL that produced it, so users can verify what was run.

Prompt Engineering and Schema Design

The system prompt is the application logic. That sounds like a metaphor; it isn’t. Every constraint, business rule, and edge-case handler that would live in code in a traditional system lives in the prompt in a production NL-to-SQL system. It deserves the same review process, versioning, and regression testing as code.

Column descriptions outperform column names by a wide margin. A column called cpa_usd is opaque. A column described as “total cost-per-acquisition in USD, including marketing attribution and fulfillment overhead, used for channel ROI calculation” is not. The model will use the latter correctly in a join it has never seen before. This is because the model is doing semantic matching between the user’s question and the schema description — the richer the description, the more signal there is to match against.

Domain-specific few-shot examples are more valuable than general SQL examples by roughly an order of magnitude. A model that has seen ten correct pharmacy fill-rate queries in its context will write the eleventh correctly under conditions where zero-shot inference would hallucinate a GROUP BY. The examples in the pool should come from real, high-rated query pairs — not hand-authored demonstrations. More on how to build that pool in the RL section.

Explicit prohibition rules prevent the most common error classes. “Never use EXTRACT without also using DATE_TRUNC” is a rule. “When the user asks about delivery performance, exclude NULL delivery_days unless they explicitly ask for in-transit orders” is a rule. These rules encode the institutional knowledge that turns a technically-correct query into a business-correct query. Below is a representative excerpt from the system prompt used in this project:

SYSTEM PROMPT (excerpt)
────────────────────────────────────────────────────

You are a read-only SQL analyst with access to the following schema.
Generate valid Amazon Redshift SQL only. Never use UPDATE, DELETE,
INSERT, DROP, or any DDL statement. If the user's question cannot
be answered from the available tables, say so explicitly.

AVAILABLE TABLES
────────────────
fills (
  fill_id          BIGINT,        -- unique prescription fill event
  patient_id       BIGINT,        -- FK → customers.patient_id (PII-masked)
  ndc_code         VARCHAR(11),   -- national drug code (non-PII)
  channel          VARCHAR(20),   -- 'prime_rx' | 'mail_order' | 'retail'
  fill_ts          TIMESTAMP,     -- UTC time of dispense
  copay_usd        NUMERIC(10,2), -- patient copay amount
  cost_usd         NUMERIC(10,2), -- total drug cost
  days_supply      INT,           -- quantity dispensed in days
  status           VARCHAR(20)    -- 'completed' | 'returned' | 'partial'
)

orders (
  order_id         BIGINT,
  fill_id          BIGINT,        -- FK → fills.fill_id
  pharmacy_id      VARCHAR(8),    -- fulfillment location code
  promised_ts      TIMESTAMP,     -- SLA promise time
  shipped_ts       TIMESTAMP,     -- actual ship time
  delivery_days    INT            -- transit days, NULL if not yet delivered
)

RULES
──────
1. Always qualify ambiguous date columns with the table name.
2. When aggregating across channels, GROUP BY channel before filtering.
3. NULL delivery_days means the order is in-transit; exclude from SLA calcs
   unless the user explicitly asks for in-transit data.
4. Use DATE_TRUNC for all period bucketing — never EXTRACT alone.
5. Prefer CTEs over subqueries for readability.

FEW-SHOT EXAMPLES (highest-rated pairs from feedback pool)
──────────────────────────────────────────────────────────
Q: "What was the 30-day fill rate by channel last month?"
A:
WITH monthly_fills AS (
  SELECT
    channel,
    COUNT(*) FILTER (WHERE days_supply = 30)  AS thirty_day_fills,
    COUNT(*)                                   AS total_fills
  FROM fills
  WHERE DATE_TRUNC('month', fill_ts) = DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month')
  GROUP BY channel
)
SELECT
  channel,
  thirty_day_fills,
  total_fills,
  ROUND(100.0 * thirty_day_fills / NULLIF(total_fills, 0), 2) AS fill_rate_pct
FROM monthly_fills
ORDER BY fill_rate_pct DESC;

RAG for Schema Retrieval

At small schema sizes, you can inject everything into the prompt. Past roughly 50 tables, you cannot — not because of token limits alone, but because precision degrades. Including irrelevant tables increases the probability of the model constructing joins that don’t exist, referencing columns from the wrong table, or conflating similarly named fields. The solution is retrieval-augmented generation applied to schema metadata.

The indexing step embeds each table and column as a natural-language description — not the DDL. “Table: fills. Tracks each completed prescription dispense event at Amazon Pharmacy. Grain: one row per fill. Key business use: fill rate, copay analysis, channel mix reporting. Joins to orders via fill_id.” That is a retrievable document. Raw DDL is not. Research on the Spider benchmark has shown that including rich foreign-key and description metadata in prompts improves execution accuracy by over 50 percentage points compared to DDL-only injection.

Graph traversal handles multi-hop joins. When the retrieved tables are fills and customers, the system needs to know whether they join directly or through an intermediate table. A pre-computed join graph — derived from foreign key constraints and manually curated where FK constraints are absent — enables the system to automatically include bridge tables when needed. Uber’s QueryGPT addresses this with a dedicated Table Agent that validates retrieved tables and resolves join paths before the final generation call.

Column pruning is the final retrieval step. Even after semantic search returns the right tables, a wide table with 80 columns will dilute the signal in the prompt. A Column Prune Agent — another of Uber’s explicit pipeline stages — scores each column for relevance to the query and injects only the top-N into the generation prompt. This reduces token usage, improves generation precision, and eliminates the risk of the model using a similarly-named but semantically different column from the same table.

Privacy First

Deploying an LLM-powered query interface on a warehouse that contains protected health information — which Amazon Pharmacy’s data does — requires a privacy architecture that goes well beyond standard access controls. The threat model has two distinct surfaces: the schema injection (what the LLM sees) and the result set (what the user sees).

The controls we implemented, in priority order:

1. PII column masking in schema metadata. Columns containing direct or quasi-identifiers — patient_id, member_id, address, date-of-birth — are tagged in the schema registry. When that column’s description is retrieved and injected into the prompt, it is replaced with a masked description that communicates join semantics (“foreign key for linking fill records to the patient record”) without describing the PII. The LLM knows the column exists and how to join through it; it cannot be prompted to surface its contents.
2. Permission-scoped schema injection. The schema retrieval step is aware of the requesting user’s IAM role and data-access entitlements. A finance analyst sees cost columns. A pharmacist sees dispensing columns. Neither sees the other’s domain unless their role explicitly grants it. The injected schema is the intersection of (semantically relevant) and (user-authorized) — not the full schema subset.
3. LLM never sees actual data. The LLM receives schema descriptions and generates a query. It never receives query results, sample rows, or any data values. The MaskSQL framework formalizes this principle — replacing sensitive schema elements with abstract symbols before they reach a remote LLM. In our system, the LLM is purely a query translator; execution and result handling are entirely outside its context.
4. Result set sanitization. Generated queries that reference PII columns directly (a rare but possible outcome of ambiguous prompting) are blocked at the result layer. A whitelist of allowed output columns — derived from the user’s access scope — filters the query before execution.
5. Audit logging. Every generation event — user ID, timestamp, natural language query, generated SQL, execution status, row count returned — is written to an immutable audit log. This is the compliance surface for any downstream data-access review.

Execution-Guided Self-Correction

The first generated query fails more often than you want. In early testing on the pharmacy schema, first-pass execution success was roughly 70% — meaning three in ten queries required either a retry or a human correction. That number is unacceptable for a tool that’s supposed to replace analyst bandwidth.

Self-correction is the practice of feeding the execution error back to the model as part of a follow-up generation call. The retry prompt includes the original question, the failed SQL, and the full error message from the warehouse. Models are remarkably good at fixing their own errors when given the error text — a type mismatch, a missing alias, a non-existent column reference. In practice, one retry pass raises first-pass-or-retry success from 70% to 88-92%. Two retry passes recover most of the remaining failures.

The failure categories that don’t recover through self-correction are almost always schema gaps — the user is asking about data that doesn’t exist in the authorized schema, or using business terminology that maps to nothing in the column descriptions. These failures are the most valuable signal in the system, because they directly indicate where the schema documentation needs to be improved.

Closing the Loop: Refinement and RL

A system that doesn’t learn from usage will plateau at its initial quality level. Reinforcement learning from human feedback closes the gap between a system that sometimes works and one that continuously improves.

The reward signal has two components. The first is implicit: execution success. A query that runs without error and returns a non-empty result is a positive signal. A query that fails, times out, or returns zero rows on a question that should have data is a negative signal. This signal is cheap — it requires no human annotation and is available for 100% of queries. It is also noisy: a query can execute successfully and still be semantically wrong.

The second is explicit: user ratings. A thumbs-up/thumbs-down on the result, with an optional correction field where the user can provide the query they intended. Explicit signals are sparse — most users don’t rate — but they are high-precision. A user-provided correction is a preference pair: (generated query, corrected query) for the same input. That is exactly the data format required for Direct Preference Optimization.

DPO is the right algorithm here over PPO-style RLHF for a practical reason: it doesn’t require training a separate reward model. DPO directly optimizes the generation model’s policy against the preference pairs using a contrastive loss. For a team without dedicated ML infrastructure, DPO can be applied as a fine-tuning pass on an open-source base model — LLaMA 3.3 in our case — using the accumulated preference pairs as training data. The result is a model that’s specialized to our schema, our query patterns, and our users’ implicit definition of a good answer.

The few-shot example pool is the faster-cycling feedback mechanism. High-rated (question, SQL) pairs — those with explicit positive feedback or a consistent execution success history — are automatically promoted into the few-shot pool. Low-rated pairs are demoted. The pool is re-ranked weekly. This doesn’t require model retraining; it improves generation quality on the next request by changing what examples are in context. For teams without the compute budget for DPO cycles, this alone will recover a significant fraction of the quality gains.

What I’d Build Next

The system as described brings first-pass-or-retry success above 90% on the bounded pharmacy analytics domain. The remaining failure surface points to three things worth building.

Ambiguity resolution as a first-class step. Instead of silently choosing a definition when the user says “active customers,” the system should detect the ambiguity and ask a targeted clarifying question before generation. This requires a lightweight pre-generation classifier that identifies under-specified terms and maps them to the business definitions that exist in the schema documentation.

Semantic result validation. Execution success is a weak correctness signal. A stronger signal comes from checking whether the result is plausible given the question — a fill rate of 340% should trigger a validation flag, not silently appear in the output. Statistical bounds on expected result ranges, derived from historical query results, would catch this class of error before it reaches the user.

Schema-aware fine-tuning on a local model. The privacy architecture currently relies on prompt-level controls to prevent PII exposure to remote APIs. A locally-hosted fine-tuned model would eliminate the remote API surface entirely — at the cost of requiring ML infrastructure to maintain. As DPO fine-tuning on LLaMA-class models becomes faster and cheaper, this tradeoff will flip for most organizations with PHI or PII data within the next eighteen months.