Data retrieval using query generation models
Systems, methods, and devices that relate to data retrieval using query generation models are disclosed. In one example aspect, the system receives a natural language request, from a user, for data stored within a database. The system can use a series of specialized large language model (LLM) agents to determine information about where the relevant data is stored within the database and how the data can be retrieved. The system then inputs this information into a query generation LLM, which generates a data retrieval query for retrieving the relevant data from the database. The system executes the query and outputs, to the user, a natural language response that includes the data requested by the user.
1 . A non-transitory, computer-readable storage medium comprising instructions recorded thereon, wherein the instructions, when executed by at least one data processor of a system, cause the system to:
receive, via a user device, a natural language input comprising a request for a data item stored in a database;
input, into a first retrieval-augmented generation (RAG) large language model (LLM), the natural language input and a listing of one or more tables within the database to cause the first RAG LLM to generate a first output identifying a table storing the data item within the database;
input, into a second RAG LLM, the natural language input and a table schema of the identified table to cause the second RAG LLM to generate a second output identifying a column of the identified table storing the data item within the table;
input, into a third RAG LLM, the natural language input and a set of Structured Query Language (SQL) queries to cause the third RAG LLM to generate a third output identifying a SQL query from the set of SQL queries,
wherein the set of SQL queries corresponds to the identified table within the database;
input, into a query generation LLM, the first output, the second output, and the third output to cause the query generation LLM to generate a final SQL query to retrieve the data item from the database;
execute the final SQL query to cause the retrieval of the data item from the database; and
output, to the user device, a natural language response comprising an indication of the retrieved data item.
2 . The non-transitory, computer-readable storage medium of claim 1 , wherein the set of SQL queries comprises pairs of: (i) a particular request for a particular data item within the table and (ii) a corresponding SQL query.
3 . The non-transitory, computer-readable storage medium of claim 2 , wherein causing the third RAG LLM to generate the third output identifying the SQL query from the set of SQL queries comprises causing the third RAG LLM to:
identify a first pair for which a first request matches the request of the natural language input; and
generate the third output identifying the SQL query, wherein the SQL query corresponds to a first corresponding SQL query of the first pair.
4 . The non-transitory, computer-readable storage medium of claim 1 , wherein the instructions further cause the system to:
generate, using a SQL query generation model, a corresponding set of SQL queries for each table within the database for retrieval of data items from each table in the database; and
convert each SQL query of each corresponding set of SQL queries into a corresponding SQL embedding.
5 . The non-transitory, computer-readable storage medium of claim 4 , wherein the instructions causing the third RAG LLM to generate the third output identifying the SQL query from the set of SQL queries further comprises instructions causing the third RAG LLM to:
convert the natural language input into an input embedding;
search the set of SQL queries for a SQL embedding that is most relevant to the input embedding; and
identify the SQL query corresponding to the SQL embedding.
6 . The non-transitory, computer-readable storage medium of claim 1 , wherein two or more of the first RAG LLM, the second RAG LLM, and the third RAG LLM are a same RAG LLM.
7 . A device comprising:
at least one hardware processor; and
at least one non-transitory memory storing instructions, which, when executed by the at least one hardware processor, cause the device to:
receive, via a user device, a natural language input comprising a request for a data item stored in a data storage location;
input, into a first generative model, the natural language input and a listing of one or more data structures within the data storage location to cause the first generative model to generate a first output identifying a data structure storing the data item within the data storage location;
input, into a second generative model, the natural language input and a data structure schema of the identified data structure to cause the second generative model to generate a second output identifying a column of the identified data structure storing the data item within the data structure;
input, into a third generative model, the natural language input and a set of Structured Query Language (SQL) queries to cause the third generative model to generate a third output identifying a SQL query from the set of SQL queries;
input, into a query generation large language model (LLM), the first output, the second output, and the third output to cause the query generation LLM to generate a final SQL query to retrieve the data item from the data storage location;
execute the final SQL query to cause the retrieval of the data item from the data storage location; and
output, to the user device, a natural language response comprising an indication of the retrieved data item.
8 . The device of claim 7 , wherein the set of SQL queries corresponds to the identified data structure within the data storage location.
9 . The device of claim 7 , wherein the set of SQL queries comprises pairs of: (i) a particular request for a particular data item within the data structure and (ii) a corresponding SQL query.
10 . The device of claim 9 , wherein causing the third generative model to generate the third output identifying the SQL query from the set of SQL queries comprises causing the third generative model to:
identify a first pair for which a first request matches the request of the natural language input; and
generate the third output identifying the SQL query, wherein the SQL query corresponds to a first corresponding SQL query of the first pair.
11 . The device of claim 7 , further comprising:
generating, using a SQL query generation model, a corresponding set of SQL queries for each data structure within the data storage location for retrieval of data items from each data structure in the data storage location; and
converting each SQL query of each corresponding set of SQL queries into a corresponding SQL embedding.
12 . The device of claim 11 , wherein causing the third generative model to generate the third output identifying the SQL query from the set of SQL queries further comprises causing the third generative model to:
convert the natural language input into an input embedding;
search the set of SQL queries for a SQL embedding that is most relevant to the input embedding; and
identify the SQL query corresponding to the SQL embedding.
13 . The device of claim 7 , wherein causing the query generation LLM to generate the final SQL query further comprises causing the query generation LLM to generate the final SQL query based on the data structure, the column, and data structure and column information included in the SQL query.
14 . A system comprising:
a communication device to:
receive, via a user device, a request for a data item stored in a data storage location; and
output, to the user device, an indication of a retrieved data item; and
a data-processing device to:
input, into a first model, the request and a listing of one or more data structures within the data storage location to cause the first model to generate a first output identifying a data structure storing the data item within the data storage location;
input, into a second model, the request and a data structure schema of the identified data structure to cause the second model to generate a second output identifying a column of the identified data structure storing the data item within the data structure;
input, into a third model, the request and a set of Structured Query Language (SQL) queries to cause the third model to generate a third output identifying a SQL query from the set of SQL queries,
wherein the set of SQL queries corresponds to the identified data structure within the data storage location;
input, into a query generation large language model (LLM), the first output, the second output, and the third output to cause the query generation LLM to generate a final SQL query to retrieve the data item from the data storage location; and
execute the final SQL query to cause the retrieval of the data item from the data storage location.
15 . The system of claim 14 , wherein the set of SQL queries corresponds to the identified data structure within the data storage location.
16 . The system of claim 14 , wherein the set of SQL queries comprises pairs of: (i) a particular request for a particular data item within the data structure and (ii) a corresponding SQL query.
17 . The system of claim 16 , wherein causing the third model to generate the third output identifying the SQL query from the set of SQL queries comprises causing the third model to:
identify a first pair for which a first request matches the request; and
generate the third output identifying the SQL query, wherein the SQL query corresponds to a first corresponding SQL query of the first pair.
18 . The system of claim 14 , wherein the data-processing device is further caused to:
generate, using a SQL query generation model, a corresponding set of SQL queries for each data structure within the data storage location for retrieval of data items from each data structure in the data storage location; and
convert each SQL query of each corresponding set of SQL queries into a corresponding SQL embedding.
19 . The system of claim 18 , wherein causing the third model to generate the third output identifying the SQL query from the set of SQL queries further comprises causing the third model to:
convert the request into an input embedding;
search the set of SQL queries for a SQL embedding that is most relevant to the input embedding; and
identify the SQL query corresponding to the SQL embedding.
20 . The system of claim 14 , wherein causing the query generation LLM to generate the final SQL query further comprises causing the query generation LLM to generate the final SQL query based on the data structure, the column, and data structure and column information included in the SQL query.