2 · The Query Pipeline

The heart of the product. src/pipeline/pipeline.py orchestrates nine stages; each stage is a module in src/pipeline/stages/.

question
   │
   ├─[1] cache        ── hit? ─────────────────────────────► return (~10ms)
   ├─[2] intent          classify: aggregation / trend / …
   ├─[3] schema select   pick ~8 relevant tables, compress to DDL
   ├─[4] sql gen         ► AI CALL — the only place NL becomes SQL
   ├─[5] sql validate    ◄ SECURITY BOUNDARY — six layers, read-only
   │        └─ invalid? feed the error back and regenerate once
   ├─[6] execute         run on the user's DB (row cap + timeout)
   ├─[7] answer gen      ► AI CALL — prose answer, streamed
   ├─[8] chart gen       deterministic, no AI
   └─[9] cache write

PipelineContext

One mutable object flows through every stage (src/pipeline/context.py). Each stage reads what it needs and writes its output.

ctx = PipelineContext(
    request_id=str(uuid.uuid4()),
    source_id=source_id,
    question=question,
    schema_data=schema_data,   # the stored snapshot
    dialect=dialect,           # drives sqlglot
    options=options,           # max_rows, include_chart, chart_type…
)
# stages fill in:
# ctx.intent, ctx.selected_tables, ctx.compressed_schema,
# ctx.generated_sql, ctx.validated_sql, ctx.query_result,
# ctx.answer, ctx.chart, ctx.tokens_used, ctx.errors

ctx.mark_stage(name) records completion — that list is what logs and job progress read.

Why a mutable context and not return values? Nine stages each returning a growing tuple is unreadable, and the streaming variant needs to expose partial state mid-flight. The trade-off is that stage order matters and is implicit; the orchestrator being one readable function is what keeps that honest.

Stage 1 — Cache

def _cache_key(source_id: str, question: str) -> str:
    normalized = question.strip().lower()
    digest = hashlib.sha256(f"{source_id}:{normalized}".encode()).hexdigest()[:24]
    return f"{CACHE_PREFIX}{digest}"

Three things to notice:

1. source_id is in the key. Two tenants asking the same question must never share an answer. There is a test for exactly this.
2. Normalisation — trim and lowercase, so “Top Products” and “top products ” hit the same entry.
3. Failures are swallowed. Both cache functions catch every exception and log a warning. A Redis outage degrades to a cache miss instead of a 500.

A hit skips all eight remaining stages: no model call, no query, ~10ms instead of ~2s. This is also what lets the public demo survive a 20-requests-per-day free tier — see Deployment.

Stage 2 — Intent

Keyword matching, no AI:

INTENT_KEYWORDS = {
    "forecast": ["forecast", "predict", "projection", "next month", …],
    "anomaly":  ["anomaly", "unusual", "outlier", "spike", …],
    "insight":  ["insight", "trend", "pattern", "why", "explain", …],
    "schema_question": ["what tables", "what columns", "schema", …],
}

Falls back to data_query. It costs nothing and is right most of the time.

Stage 3 — Schema selection

This is the retrieval problem, and it is the part most text-to-SQL demos skip. Covered in depth in Schema & Connectors; the summary:

1. Score every table against the question — table-name matches weigh 3×, column names 1×, sample values 0.5×, plus a bonus for large tables (fact tables are usually what questions are about).
2. Follow foreign keys. “Revenue by product” scores orders and products but needs order_items to join them — and nothing in the question mentions it.
3. Compress the selection into minimal DDL.

Selection is deterministic by construction — ordered by score, then name, with FK additions appended sorted. It used to return list(top_set), and Python randomises string hashing per process, so the same question produced different prompts across restarts. See Bugs That Shipped.

Stage 4 — SQL generation

The one place natural language becomes SQL.

system_prompt = template.render(
    dialect=ctx.dialect,
    schema_ddl=ctx.compressed_schema or "",
    max_rows=ctx.options.get("max_rows", 1000),
)
result = await provider.complete(
    messages=[Message("system", system_prompt), Message("user", ctx.question)],
    options=CompletionOptions(temperature=0.0, max_tokens=1024),
)

Temperature 0.0. This is translation with a right answer, not creative writing. Say that if asked why.

Models fence output inconsistently even when told not to, so _clean_sql unwraps it:

_FENCE_BLOCK = re.compile(r"```(?:sql|postgresql|mysql)?\s*(.*?)\s*```",
                          re.DOTALL | re.IGNORECASE)

It prefers a complete fenced block, falls back to stripping whichever fence is present, and removes a trailing semicolon — because connectors wrap SQL as SELECT * FROM (<sql>) _q LIMIT n, which a semicolon breaks. The original implementation sliced off the first and last line, which mangled unfenced SQL.

Stage 5 — SQL validation

The security boundary. Full treatment in Security. Six layers:

The repair loop

MAX_SQL_ATTEMPTS = 2
 
async def _generate_valid_sql(ctx, provider) -> bool:
    original_question = ctx.question
    try:
        for attempt in range(MAX_SQL_ATTEMPTS):
            sql = await generate_sql(ctx, provider)
            if not sql:
                ctx.add_error("SQL generation returned an empty result.")
                return False
            if await validate_and_set(ctx):
                return True
            if attempt < MAX_SQL_ATTEMPTS - 1:
                error_hint = "; ".join(ctx.errors)
                ctx.errors = []
                ctx.question = (
                    f"{original_question}\n\n"
                    f"A previous attempt produced invalid SQL and failed with: {error_hint}\n"
                    f"Return corrected SQL."
                )
        return False
    finally:
        ctx.question = original_question

Three details worth pointing at:

• Bounded. Two attempts. Unbounded self-correction burns tokens on a hopeless question.
• Errors cleared before the retry, so a successful repair does not report as “partial”.
• finally: ctx.question = original_question. The repair prompt is written into the context. If it were left there it would reach the answer prompt and — worse — the cache key, so a repaired query would never hit cache again. There is a test for this.

Stage 6 — Execute

result = await connector.execute_query(sql=sql, max_rows=max_rows, timeout=timeout)

The connector wraps the statement:

limited_sql = f"SELECT * FROM ({sql}) _q LIMIT {max_rows}"

So a model that ignores the row limit in its own SQL still cannot materialise a hundred-million-row result. Failures are captured onto ctx.errors rather than raised, so the caller gets a partial result with an explanation instead of a 500.

Stage 7 — Answer generation

Rows go back to the model to be described in prose, at temperature 0.3 (slightly higher — this is writing). Two entry points: generate_answer() for sync, stream_answer() yielding tokens for SSE.

Stage 8 — Chart generation

Deterministic. No AI call. Chart type is inferred from result shape:

Stage 9 — Cache write

# Awaited rather than fire-and-forget: a bare asyncio.create_task is only
# weakly referenced by the event loop and can be garbage collected before it runs.
await write_cache(ctx, ttl=settings.query_cache_ttl_seconds)

That comment is a real bug fixed. asyncio.create_task returns a task the loop holds weakly — without a strong reference it can be collected mid-flight, so cache writes silently vanished under load. Same trap appears in src/workers/dispatch.py, which keeps a module-level set of in-flight tasks.

Streaming

stream_pipeline runs the same stages but yields (event, ctx) tuples:

thinking → thinking → sql → thinking → token × N → complete → done

The UI shows the SQL as soon as sql arrives, long before the answer finishes. It yields the context alongside each event so the service can log the completed query after the stream closes.

Failure modes, and what the caller sees

That last row is a bug that reached production: /v1/query answered 500 while /v1/query/stream answered correctly, because only the streaming path guarded its logging call. See Bugs That Shipped.

Read the tests

tests/test_pipeline.py drives all of this against a real SQLite database with only the model faked. It is the best single file for understanding pipeline behaviour, because each test names a property rather than an implementation detail:

• the repair loop stops after one retry
• the repair prompt never leaks into the answer or the cache key
• DROP TABLE is rejected by the pipeline, not just by the validator’s own unit tests
• a failing query returns a usable context instead of raising