Build an Agent to create and run queries for users

SQL Generation with Generative AI

OVERVIEW • AGENT TOOLS • CREATE AND TEST AGENT • CONCLUSIONS

Users who want to get information from databases used to require a good understanding / knowledge of SQL.

In addition to this, they needed to use one or more applications to access the data.

With advances in Generative AI, we are able to understand user text input and generate different types of content, including code

In this demo, we will create an Agent user’s can interact with to get information from databases without SQL knowledge.

We will use Google Generative AI, Langchain, Bigquery to build our solution

Overview

{...}/<...>

SQL

Process Diagram

Embeddings

Textembedding-gecko

SGL Generator

Gemini Pro

Debugging Agent

Code Chat

Report Comparison

Gemini Pro

Visualization

Generator

Gemini Pro

Marketing

Recommendations

Gemini Pro

Source DB

BigQuery

Vector DB

Known Good SQL

Table Embeddings

Check Known Good SQL DB

Infer Schema

Generative SQL

Validate SQL Query

(BQ Dry Run)

Execute Query

Generate Visuals

Final Response

• Samples of good Known sql queries are collected. They will serve for few-shot prompting when generating SQL code.
• Tables descriptions and schemas are passed to an LLM to get auto descriptions
• Both items are converted to embeddings
• Vector DB index is build
• Index is deployed to an endpoint to serve for predictions

Process Description: Set up Vector DB

Vector DB

Known Good SQL

Table Embeddings

Embeddings

Textembedding-gecko

Agents and Tools Overview

Interact with the user.

Decides with action to take from a sequence of actions available.

Can use LLM as the reasoning engine to determine next action.

Agents

Functions that an agent can use to interact with the world

Main considerations when defining agent’s tools: Give access to the right tools, Describe properly to be helpful for the agent.

There are many ways to define tools. In this demo, we will use the @tool decorator

Tools

Before we start creating the Tools we may need to create some util functions that will help us to perform the actions desired.

Utils functions that we will use in this demo: Get samples from Vector DB, Get a code chat, Generate a SQL

Tools

• The process starts with the user query
• The query is converted to embeddings and matched with good known queries (if match)
• If they match then we just need to run the query. Otherwise we need to proceed generating it
• To generate the query we get the relevant tables and examples from the vector DB
• Then, a prompt is sent to Gemini to get a query candidate
• Big query Dry run is used to check for errors
• A code chat is used to fix the query if errors found. This is an iterative process
• Once there is a query without errors it run on BQ to get the data
• The results of the data and additional context is passed to Gemini to generate python code to be executed to generate a visualization.
• If the execution is successful the visual is included as a base64 image in the response
• Finally, the agent generate a final text answer based on the data and the original question.

Agents and Tools Overview

Tool Example:

generate_executive_report

Run necessary tasks to provide to generate an executive report

Tasks involved:

Get samples and relevant tables

Iterate until getting a good query

Run query

Use data results to generate a report

Get Relevant tables /queries

RAG

Generate SQL

Code Chat/Gemini

Run Query

Big Query

Generate Report

Gemini

Now that we have the set of tools we can create the agent.

We will provide to the agent a memory so to allow a conversation.

We can enable verbose in the agent in order to see in detail the workflow of decisions

There are many type of agents we can choose:

Initialize and Test Agent

• Tool calling
• OpenAI Tools
• OpenAI Functions
• XML
• Structured Chat
• JSON Chat
• ReAct
• Self Ask With Search

With Generative AI we can automate the process of generating queries for user with no sql knowledge

We can combine Text Generation and Code Generation models into a single agent to provide with the tools required to perform many actions and access more complex tools.

Using RAG for retrieving relevant tables and query samples improves significantly the performance of the agent.

Key Takeaways