Open Source AI Database Agent Part 1: Introduction

Generative AI is top of mind for many engineers. The questions of how it can be applied to solve business problems and boost productivity are still up in the air. Recently I wrote a blog post about the impact of AI on platform engineers, where I talked about various AI Agents and how they can interact with different systems and APIs. The belief is that AI Agents or a combination of those would heavily impact how we work with computers and systems today.

One of the examples that I want to look into today is a Database AI Agent. We may call it a text-to-SQL model that translates human-like requests into SQL language and produces human-like responses.

We will use a combination of open source tools:

• Langchain – An open source python library that offers developers a comprehensive set of resources to develop applications that run on Large Language Models (LLMs). 
• Percona Distribution for PostgreSQL—We will use the Distribution to host the database with sample data. I will use Pagila, the famous Sakila example database ported to PostgreSQL. The structure of the database can be found here.

All the scripts from this blog post are available in this GitHub repository.

This blog post series has two parts:

1. Open Source AI Database Agent Part 1: Introduction (this one) – Covers the basics, we create a simple database AI agent.
2. Open Source AI Database Agent Part 2: OSADA—We fine-tuned the scripts, expanded its capabilities, and added the graphical UI.

The setup

Provision the database

Install Percona Distribution for PostgreSQL

I’m running Ubuntu, so I will follow our Quickstart guide for apt repositories. 

• Install percona-release

• Setup the apt repo for PostgreSQL 16

• Install the distribution

Bootstrap the database

Get Pagila database SQLs from GitHub:

Create the database:

Create the schema:

Insert the data:

I will also create the user pagila to access the database from my scripts:

sudo su postgres

Setting up Langchain

With Langchain, any LLM can be used. In our examples, we will use Google Gemini models, which require a key. You can also use OpenAI in the next iteration of the agent; read more in part two of this series..

Step 0 – Make sure Gemini works

Start with installing the required Python packages. The usual recommendation is to use pip:

A simple script, 0_langchain_gemini_test.py, asks Gemini, “What can you tell me about Percona?” You can set the GOOGLE_API_KEY variable, but if not, the script will ask for it. 

It is working.

Query the database

We are done with the initial setup; let’s query the database.

Install the required libraries:

To experiment with querying the database we use 1_langchain_gemini_postgresql.py script. The script takes a user question in a human-readable format as an argument and produces the response.

To do that we will use the create_sql_agent function that constructs an agent specifically designed to interact with SQL databases. This agent combines a language model (like Gemini) with specialized tools that understand and manipulate SQL queries.

In addition to GOOGLE_API_KEY, we need to provide details about the database:

Make sure that you set PG_USERNAME  and PG_PASSWORD  as environment variables. 

Let’s start with the simple question:

Without providing any additional context or explaining the structure of the database, I was able to get a proper answer. The way it works is the following:

But this works only for simple questions, where it is quite straightforward. For more complex questions, multiple requests are required for the database and LLMs, so-called chains. It is also quite interesting, that chains help to fix some errors in the query on the fly. Below, you can see how the text-to-sql script worked for two more questions:

1. Which 5 actors have appeared in the most films?
2. What are the most popular film categories? Show the top 3 with the number of films in that category.

Known issues

Sometimes, the Large Language Model can’t fix the error in the SQL query, and the agent gets stuck in the loop. It can happen for the same question that was successfully answered before. In the case of the loop, the agent will notify you about it:

The potential impact

Companies generate tons of data that they later use to make better business decisions. A perfectly structured BI system requires the skills of analysts and engineers. In this article, we saw how generative AI can help digest the data from rudimentary databases, but the expectation is that it will mature over time. There will be no need to learn SQL or understand how data is stored; AI will make sense of it automagically. This is going to impact how businesses store, manage, and process data.

In the second part of the blog post, we productize the script:

1. Add UI and a docker container for easier onboarding
2. Add MySQL to supported databases
3. Add OpenAI and custom LLMs

Check out Open Source AI Database Agent Part 2: OSADA!