Imperva uses Amazon Bedrock to optimize SQL generation from natural language

This is a guest post co-written with Ori Nakar of Imperva.

Imperva Cloud WAF protects hundreds of thousands of websites from cyber threats and blocks billions of security events daily. Counters and insights based on security events are calculated daily and used by users across multiple departments. Millions of counters are added daily and 20 million insights are updated daily to find threat patterns.

Our goal was to improve the user experience of an existing application used to explore counter and insight data. The data is stored in a data lake and retrieved by SQL using Amazon Athena.

As part of the solution, we replaced multiple search fields with a single free text field, using a large language model (LLM) with example queries to enable searches to be performed using the language used by Imperva's internal users (business analysts).

The following diagram shows the search query that was translated into SQL and executed. The results were later formatted as graphs by the application. There are various types of insights like global, industry and customer level insights that are used by multiple departments like marketing, support, research etc. The data was made available to the users through a simplified user experience powered by LLM.

Figure 1: Natural language insight search

Amazon Bedrock is a fully managed service that offers a choice of high-performance foundational models (FMs) from leading artificial intelligence (AI) companies, including AI21 Labs, Anthropic, Cohere, Meta, Mistral, Stability AI, and Amazon, within a single API. It also provides a wide range of capabilities needed to build generative AI applications with security, privacy, and responsible AI. Amazon Bedrock Studio is a new single sign-on (SSO) enabled web interface that provides a way for developers across an organization to try out LLMs and other FMs, collaborate on projects, and iterate on generative AI applications. It provides a rapid prototyping environment and streamlines access to multiple FMs and developer tools within Amazon Bedrock.

You can read more about this issue and how we used Amazon Bedrock to achieve high-quality results in our experiments and deployments here.

problem

Giving users access to data through applications has always been a challenge. Data is usually stored in a database and can be queried using the most common query language, SQL. Applications use various UI components to allow users to filter and query the data. There are applications with dozens of different filters and other options, all built to make data accessible.

Querying a database through an application is not as flexible as running SQL queries with a known schema. To empower users further, a simple user experience (UX) is required. Natural language solves this problem by supporting complex, yet easy-to-read natural language queries without any SQL knowledge. If the schema changes, the application's UX and code remain the same or require only minor changes, saving development time and making the application's user interface (UI) stable for users.

Constructing SQL queries from natural language is not an easy task. SQL queries must be syntactically and logically correct. Using LLM with proper examples makes this task easier.

Figure 2: High-level database access using LLM flows

Challenge

LLM can build SQL queries based on natural language. The challenge is to ensure quality. Users can enter any text and the application builds queries based on it. There is no option to cover all options and ensure the application works correctly, as in traditional applications. Adding LLM to an application adds an extra layer of complexity. Responses by LLM are not deterministic. Examples sent to LLM are based on database data, making it more difficult to control requests sent to LLM and ensure their quality.

The solution: A data science approach

In data science, it is common to develop a model and then fine-tune it using experiments. The idea is to use metrics to compare the experiments as they are developed. Experiments can differ from each other in many ways, such as the inputs sent to the model, the type of model, or other parameters. The ability to compare different experiments enables progress; you can learn how each change contributes to the model.

The test set is a static set of records that contains the prediction result for each record. When predictions are run on the test set, the results are recorded along with the metrics needed to compare the experiments. A common metric is precision, which is the percentage of correct results.

In our case, the results generated by LLM are SQL statements. SQL statements generated by LLM are not deterministic and are difficult to measure, but running the SQL statements on a static test database makes them deterministic and measurable. As a test set, we used a test database and a list of questions with known answers. This allowed us to run experiments and fine-tune our LLM-based application.

Database Access with LLM: From Question to Answer

Given a question, we defined the following flow: The question is sent to the Search Extension Generation (RAG) process which searches for similar documents. Each document contains a sample question and its information. The relevant documents are created as prompts and sent to the LLM which creates the SQL statements. This flow is used both in development and at application runtime.

Figure 3: Flow from question to answer

As an example, consider a database schema with two tables, orders and items. The following diagram shows the questions for the example SQL flow:

Figure 4: Example of a flow from question to answer

Database Access Using LLM: Development Process

To develop and fine-tune the application, we created the following data sets:

• Static test database: Contains sample copies of relevant tables and data.
• Test Set: Contains the questions and the resulting answers from the test database.
• Question to SQL Examples: A set of questions and their translations to SQL, some of which include the data returned, allowing you to ask questions about the data and not just the schema.

The development of the application is done by adding new questions and updating different datasets, as shown in the following diagram.

Figure 5: Adding a new question

Updates to the dataset and other parameters are tracked as part of adding new questions and fine-tuning the application. We used a tracking tool to track information about our experiments, such as:

• Parameters such as number of questions, number of examples, LLM type, RAG search method, etc.
• Metrics such as accuracy and SQL error rate
• Artifacts such as a list of erroneous results including generated SQL, returned data, etc.

Figure 6: Experimental flow

Using the tracking tool allowed us to make progress by comparing experiments. The following chart shows the accuracy and error rate metrics for the different experiments we ran.

Figure 7: Accuracy and error rate over time

If a mistake or error occurs, drill down into the erroneous results and experiment details to understand the cause of the error and correct it.

Experiment and deploy with Amazon Bedrock

Amazon Bedrock is a managed service that offers a choice of high-performance foundation models so you can try and evaluate the FM that best suits your use case and customize it for your data.

Amazon Bedrock makes it easy to switch between models and embedding options. Below is some example code using the LangChain Python library that allows you to use different models and embeddings.

import boto3
from langchain_community.llms.bedrock import Bedrock
from langchain_community.embeddings import BedrockEmbeddings

def get_llm(model_id: str, args: dict):
   return Bedrock(model_id=model_id,
                  model_kwargs=args,
                  client=boto3.client("bedrock-runtime"))

def get_embeddings(model_id: str):
   return BedrockEmbeddings(model_id=model_id, 
                            client=boto3.client("bedrock-runtime"))

Conclusion

We used the same approach used in data science projects to construct SQL queries from natural language. The solution presented here can be applied to other LLM-based applications, not just SQL construction. For example, it can be used for API access, constructing JSON data, etc. The key is to use experimentation to create test sets with measurable results and progress.

With Amazon Bedrock, you can use and switch between different models to find the one that suits your use case. You can compare different models, including small models, for better performance and cost. Because Amazon Bedrock is serverless, you don't need to manage any infrastructure. You were able to quickly test multiple models and finally integrate and deploy generative AI capabilities into your application.

You can start experimenting with natural language to SQL translation by running the code samples in this GitHub repository. This workshop is divided into multiple modules where each module builds on the previous one while introducing new techniques to solve this problem. Many of these approaches build on existing work in the community and are cited accordingly.

About the Author

Olli Nakar He is a lead cybersecurity researcher, data engineer, and data scientist in the Imperva Threat Research group.

Eitan Serra is a Generative AI & Machine Learning Specialist Solutions Architect at AWS. He works with AWS customers to provide guidance and technical support, helping them build and operate Generative AI & Machine Learning solutions on AWS. In his spare time, he enjoys jogging and reading the latest machine learning articles.

Elad Eisner I am a Solutions Architect at Amazon Web Services, working with AWS enterprise customers to design and build solutions in the cloud and help them achieve their goals.