Introduction to NL2SQL with LlamaIndex
Hi everyone. In today's video, we're going to explore the natural language to SQL capabilities of the Llama index. We have already covered, in one of the videos, explaining the basics of the Llama index and how to create, let's say, question answering or the RAG kind of app using the Llama index.
And before that, I had also covered the LangChain SQL agent, that how do you connect your own SQL database using natural language and the LangChain SQL agent. But recently, Llama index has been doing really good progress with a lot of features, and I thought this is one of the good things that we could explore because I have been using this natural language to SQL from the Llama index in a couple of my recent projects.
So in this video, we're going to see how we can connect to the database using Llama index, ask queries in the natural language, similar to what we have seen for the LangChain SQL agent. So the first thing, we're going to install the Llama index and the MySQL because when we want to connect to the database, I'm using is actually the MySQL, so we need a MySQL. So let me install those two things. And then I'm adding some basic logging capabilities so that we could see what's happening under the hood, and for the Llama index specifics, those are things I'm logging.
Database Connection Setup
So here are my database details. My database is hosted in the RDS, Amazon Relational Database Service. Whatever your database will be, it will be having its username, password, the host, which is the URL of your database, and maybe the database name. So I'm using this classic model, this is the same database I have used in the earlier example because it has a good table and the familiarity. It has like an e-commerce kind of thing, like it has customers, products, orders, payments kind of information. Something a kind of a sample DB that I could use for tutorials like this. So maybe I can put that link also. I have put it in the earlier because I'm going to share this notebook with you, right? So I can put the database link also here.
So how does this Llama index work? If you go to the Llama index, it has some kind of a wrapper class on your database. It's called the SQLDatabase, which is going to take an engine. That engine is nothing but the SQLAlchemy engine. So first, we need to create the SQLAlchemy engine that this SQL data wrapper can use to connect with our database. We can simply from SQLAlchemy, we can simply import the create_engine.
So first thing we need to create is the connection string. A connection string is going to have all these details, all the authentication URL details that will be used to connect with the thing, right? For MySQL, the connection string looks like the database which is MySQL, the driver or the library I'm using to actually connect with that particular database, and then username, password, host, and the database name. If you are using something else, let's say PostgreSQL, then make sure you know what is the connection string of PostgreSQL using, let's say, SQLAlchemy, and then you can substitute that thing here. And then once you have the connection string, you simply use the create_engine function and get the engine. And to know whether we got the correct functionality working or not, we can simply use that engine to connect to our database and fire some simple query, like I'm selecting some customers from there, to make sure that the engine or the connection string is correct and we are able to connect.
Okay, maybe I haven't run this cell. Let me run that cell and then we could run this cell. Okay, so seems to be whatever the connection string and the engine I created, I'm able to connect with my own database, and that's why I'm able to display this results, right?
So here is a table details dictionary I have. Here I have the name of the table and maybe just one-line description of the table. This is something an optional thing because sometimes what happens is the database table names what you have, those table names might not be self-explanatory. Maybe you have a customer table, but then you are calling it XYZ. So to better to have some description about your table that we can give it to, let's say, Llama index or the Chat GPT so that it understands your tables better.
So I just created one simple dictionary for the mapping purpose: table name and the description. Now, as we saw here, right, we need to create the object of the SQLDatabase. So we import the SQLDatabase from the Llama index, and this SQLDatabase will take the engine that we created above, the SQLAlchemy engine, and it can also take certain extra parameters. So it can take like this also. For example, when we are creating the SQLDatabase object, we can literally specify the tables that it should include. Let's say I have many tables, but I only want to include this table. For this, we could specify those tables. We can also... so whenever this SQLDatabase object gets created, it will get the information about the schema and it can also fetch certain records from each table so that it understands the nature of the data. And you can specify, you know, sample rows in each table kind of a thing, five, by default is three actually, but you know.
So now what I'm going to do, I'm going to allow it to take all the tables what I have, and maybe sample rows from each table is just two rows. Let's see what happens. So it will connect to the database and get... You can see what... a a SQL database object. We can check the all tables that it is able to retrieve from the database. And you could see it has got access to this many tables. You can also check the table info where you can see what particular info it got from each table. So you can see, first of all, it got the schema of the customers table, and then it also got the two rows from that customer table. So this is the information it can provide to Chat GPT so that Chat GPT understands the columns it has and the nature of the data it has. This value, we can change two rows, zero rows, or whatever, five rows, right, we want to pass. Then it will have another table called employee. Okay, this is how it is able to extract the schema and the data information that it can use actually further.
Configuring the LlamaIndex Service Context
We are also going to require actually the OpenAI key because ultimately it's going to use a large language model like OpenAI under the hood to create natural language to SQL. So I'm just putting my key into the environment variable here. So this key I will be deleting after recording this particular video so you won't be able to use this particular key, right?
Now, let's see what extra thing we are doing. So since we want to experiment with the different things, I'm just using this particular utility from the Llama index called TokenCountingHandler. This is some kind of a callback that we can add so that we know how much token we consume. This is not something needed for this natural language to SQL, but it is good to have so that I know how much... You can read more about it again.
So how does it work? First of all, we need to import this TokenCountingHandler. We give some tokenizer information and also the model that we're going to use, which is GPT 3.5 Turbo. And then we attach this token counter through the callback manager, and this callback manager can be passed inside the service context. So in the Llama index service context, as per my understanding, is something which takes some configuration, and this configuration you can use, right? So this service context... so here what configuration we have added here, you could see we have also imported OpenAI from the Llama index and actually we are creating 3.5 model as an LLM. So the service context has two information: what LLM we are using and what callback manager or any extra things we want to add in this service manager. So the service context for example, let's say you are chunking your documents in Llama index, let's say while creating the index. So in the service context, you can also specify the chunk size and all. So this is nothing but just kind of providing the configuration kind of a thing. Okay, so we got the service context, this configuration.
Basic Querying with NLSQLTableQueryEngine
Now this is the time to actually create the natural language to SQL query engine. So again, in Llama index, the query engine is actually, you can think of the query engine as a generic interface where you can ask a question to the data. That's it. You can ask a natural language question to your data, and that's what is nothing but the query engine. And then you can have a very specialized query engine. Now this query engine is specifically for the natural language to SQL table related query engine, so we call it SQLTableQueryEngine.
Now this query engine requires your database object that we created and the additional configuration and information. And then once we get this query engine... again I forgot to run the upstairs cell... okay, there's a couple of cells I need to run. Let's run this cell, which is nothing but this token counting Handler. This is like defining our configuration of the service context, and eventually, finally, our query engine, which takes all this information. Now we can ask the query.
We can ask something like, you know, maybe which customer has the highest orders? And we can execute that particular query. And you could see under... this is what we have added the logging. Now you could see the moment we ask any query, the first thing it is doing is actually retrieving or using all the schema it has. So it has nine tables, it is going to put that all nine table information here and give it to GPT-3.5 so that GPT-3.5 can understand from each table which columns to use. We got the response back and let's see what response, so we can see exactly what response we got. So the customer with the highest number of orders is this particular, whatever the name of the customer is given, you know, and they have 26 orders. Then you can look at when this response was generated, what was the SQL query actually generated internally. So it generated internally this particular query. And you can also see when it fired this query against the database, what result it got back. So this is what it got back from the database. So it used this information and eventually synthesized these answers so that we can give it to the user. Pretty simple, right?
We just simply use some kind of a wrapper called NaturalLanguageSQLTableQueryEngine, which took the database connection object and then some extra configuration detail, right? We could fire the other query and we should get, let's say, we are asking this information: when was this particular order was shipped? Again, it is going to use all this tables information, generate the response, and give it to... So it is saying it is from the 10th January, something particular. And this is what it got from the database. From the database it got the date and it did form the query. Now okay, I should have printed this thing. So this will not be the correct thing. Let's reset the counter and then run actually for one single query and see actually how much it has consumed. Okay, so it consumed almost 900 tokens when we ask any question.
The Problem with Sending the Full Schema
Now let's reset the counter because I'm going to use a different technique and then we're going to see how much... Now, if you observe what we have been doing, every time we ask questions, we are giving the whole schema to Chat GPT to figure out the column names and the tables it should use. But here we have only a couple of tables, but what if you have hundreds of tables? You can't put the hundreds of tables' information inside this every time the prompt is going to run. So you need some way so that you first figure out out of all the tables you have, what subset of the tables, maybe three or five subset of the table that you should pass to the GPT, so that it reduces your cost and maybe feasibility also. You can't put all this information in a prompt, you will get this token limit error. So how do we do that?
Building an Index on Table Schemas
So in the Llama index, what we could do, we can actually create the index on this information what you see. So for each table, we have the table schema and its description. We can create first an index on this particular information. And then before we directly query to Chat GPT, we will query this index and we will see out of this many tables, what two or three tables maybe I require to pass to the GPT. And we're going to see that particular thing now.
So to do that thing, we are going to create actually an index over the tables. So we require two things. That kind of index is called actually the Objectndex in Llama index. So Objectndex means you actually want to put some kind of objects, like in our case we're going to put now a schema object or the table object inside this Objectndex. We require two things: one is SQLTableSchema, which is nothing but it will take your table schema and give you some kind of a wrapper or some kind of interface, and SQLTableNodeMapping, which is nothing but going to... so I will tell you where this mapping is going to be used.
So first thing, let's do this thing. We already have a list of tables that we have read here, and we give this mapping, and sorry, this list of tables what we have in this SQLDatabase object, and it will create the mapping. What is the SQLTableNodeMapping? In Llama index, whenever you create the index, your data is actually get stored in an internal representation called a node. So one node is actually one chunk, you know, at a high level in your index. So here what it will do, here one node will be one SQL table schema. So this Objectndex will be having multiple such nodes and each of these nodes will have some table schema. That's what we are doing here. So first of all, we're going to iterate through all the tables and we're going to create the instance of SQLTableSchema for each table, and we're going to also add some extra contextual information which is nothing but what I have declared on the top here: store the customer data, store the list of sale number, right? This is the way you can pass some extra information to GPT, or let's say that index, that will help to narrow down your search. So come back here. So we create... so what we're going to do, we're going to collect all the schema objects in this particular list. So we iterate through our tables what we already have, and then for each table we create this SQL table schema where we give the name of the table and we also give the description what we have.
Now, once we got both the table node mapping and the schema list, we are going to create the Objectndex. So Objectndex is going to consume this table schema objects. The table_node_mapping is nothing but the table name and the node mapping, and this is what it is going to create: the vector store index. We are going to also pass the service context which talks about which LLM to use and our counter it's going to use, right? So let's first of all create this thing, table schema objects... Table details is not defined. Let me check the spelling. Okay, so we haven't defined, so let's run this table details. Come back here and now let's check again once again. Okay, so it is created. So you can see the table schema object has multiple table schemas. Okay, for each table, you can see we have this schema object where it indicates the name of the table and some extra configuration, that extra information that we have added. This is optional, this extra information is an optional thing. Okay, now let's create the index. So we are going to create an index using that same mapping and the scheme objects. So we are going to create the vector store Index, this particular index... okay maybe we first create the index, take this lines out of it. Let's create the index.
Advanced Querying with a Retriever Engine
Now, once we have this index what we created, we can use something called SQLTableRetrieverQueryEngine. So what was we earlier using? We were using NaturalLanguageSQLTableQueryEngine. Now this is like a retriever query engine. Why retriever? Now it's actually going to first of all retrieve the tables that it should use, and then eventually find the SQL corresponding query. So whenever we ask something, let's say how many customers we have, a simple thing, what it's going to do actually, you see it's actually not taking all the tables we have. It has only taken I guess two or three tables. It can take top k is equal to three tables. I think it took customers table, orders table, and the payments table. Okay, and then maybe only one of this table is actually sufficient to answer. And it has consumed 500 tokens. So this is a better way actually.
I did ask a simple question, maybe you can ask some better question, which needs to be careful. What if the table requires actually answering from multiple five tables or something? So you can put here five, six, whatever the number you want to put. But the whole idea is we are doing this thing because there is an object index retriever. This retriever is going to take top three tables from this particular Objectndex. And we have created this index because we can't put... because let's say we have hundreds of tables or maybe 50 tables and we can't put all 50 tables in each prompt request. That is the reason the first step is happening is actually retrieving the relevant tables and then actually using those retrieved tables to generate the answer. And you could see it will be consuming less tokens compared to what we were consuming.
Conclusion and Next Steps
So this was like a short tutorial to give you the glimpse of the natural language to SQL capabilities of the Llama index, and hopefully you find it useful. If you haven't seen the earlier Llama index videos, you can go and check. So we had this Llama index basic introductory video and then we also have a LangChain SQL agent, and then you can compare what we have with the LangChain SQL agent with the Llama index SQL agent. Okay, thank you.