My first text-to-SQL RAG pipeline | Notion
Author's note: This post was adapted from a presentation at the Recurse Center. If you prefer to explore hands-on first, the interactive app can be found here.
Suppose we have a text-to-SQL agent [1] and the Titanic dataset in our relational database [2]. To the input text "Number of passengers", a reasonable response from the agent is something like "SELECT COUNT(*) FROM passengers".
If the table name is dim_passengers, however, the query should be "SELECT COUNT(*) FROM dim_passengers". Let’s take this a step further. What if we want only the passengers above the age of 21? How does the agent know the column to filter on is age, passenger_age, or something else?
To provide the necessary context to the agent, we’ll need to add a separate step to retrieve that context, and augment the generation by sharing that context alongside the input text. This is commonly referred to as retrieval-augmented generation, or RAG.
First, we load up the vector database with the table metadata. The table metadata for dim_passengers could look like this (full table metadata and sample rows here).
