ChatGPT Action to query database

In this article, we will show how to integrate ChatGPT with Postgres database. We will create Python notebook in Jupyter Lab (opens in a new tab). It will have Text widget in which we can pass SQL query. The notebook will execute query in the database and displays the response as Markdown table. The notebook will be deployed as PRIVATE site with Mercury Cloud service. The ChatGPT will talk with notebook using API Key from Mercury Cloud.

Python code

We need to install few packages to make this integration works:

# install package for loading secrets from .env
pip install python-dotenv
 
# install postgres connector
pip install psycopg2-binary
 
# install widgets framework
pip install mercury 
 
# install packages to display results as Markdown table
pip install pandas tabulate

Let's import required Python modules in the first code cell:

import os
import psycopg2
import pandas as pd
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv
import mercury as mr

We keep database credentials in the .env file. We need to load them. Please read more about how to store secrets in Jupyter notebook (opens in a new tab).

_ = load_dotenv()

Let's define a function with database connection details:

def db_engine():
    host = os.environ.get("DB_HOST")
    port = os.environ.get("DB_PORT")
    user = os.environ.get("DB_USERNAME")
    password = os.environ.get("DB_PASSWORD")
    db = os.environ.get("DB_DATABASE")
    return f"user='{user}' password='{password}' host='{host}' port='{port}' dbname='{db}'"

The widget where we can write SQL query:

sql = mr.Text(label="SQL query", value="select 1", sanitize=False, url_key="query")

Please notice, that we set url_key in the constructor, it will be used in the OpenAPI schema as parameter name.

Execute SQL query in Postgres database adn fetch results:

try:
    with  psycopg2.connect(db_engine()) as conn:
        with  conn.cursor(cursor_factory=RealDictCursor) as cur:
            cur.execute(sql.value)             
            rows = cur.fetchall()
except (Exception, psycopg2.DatabaseError) as error:
    print(error)  

Get results to Markdown table:

md = pd.DataFrame(rows).to_markdown()

Display as Markdown in the notebook:

mr.Md(md)

Set result in APIResponse:

response = mr.APIResponse({"result": md})

The full Python code in the notebook:

Deploy notebook

We need to make our notebook available online. Thanks to Mercury framework (opens in a new tab) the notebook can be served as web appliction without any changes in the notebook code. Mercury is open-source solution, and can be self-hosted. However, the easiest way is to deploy notebook online is to use our hosted service called Mercury Cloud. Below are described steps to deploy notebook:

Login to Mercury Cloud

If you don't have account, you can create it here: Mercury Cloud (opens in a new tab).

Create new site

Please create a new site. Remember, to set site as PRIVATE otherwise, anyone with link to your site will be able to query your database.

Upload file with notebook

Please upload your notebook file and requirements.txt file with all required packages listed.

Provide connections secrets

Please add database credentials as secrets in Mercury Cloud.

Configure ChatGPT builder

The Mercury automatically generates OpenAPI schema, please find it by clicking the OpenAPI link in your site footer. The schema should be pasted in the GPT builder as new Action:

Please also set Authentication. Please set it to API Key and type to Bearer:

Your API Key for Mercury Cloud you can find in https://cloud.runmercury.com/account:

OK, we are ready to send some queries to our Postgres database:

The Mercury is using long polling for executing notebooks, which means that you send request to process the notebook, get the task_id as response, and then query the server with task_id till you get the final response. Below is the example with 3 requests to get the final response:

OK, let's try to run some more difficult query:

Summary

You can build custom ChatGPT that will query your database. You can use Mercury framework for build interactive notebook and serve it online. The Mercury Cloud provides you the easiest way to put your notebook online. In this article, we created a read-only user in database, so ChatGPT can only query database to show results, it can't update database. What is more, it is good to provide description of available tables, otherwise ChatGPT will try to guess table name, which can cause halucinations.

Good luck!