Chat with your holiday song database with Neon Postgres and LangChain

mkdir holiday_songs_app
cd holiday_songs_app

python -m venv venv

pip install LangChain LangChain-experimental openai python-environ psycopg2-binary

postgresql://daniel:************@ep-snowy-lake-123456.us-east-2.aws.neon.tech/holiday_songs?sslmode=require

OPENAI_API_KEY=[your_openai_api_key]
DATABASE_URL="postgresql://daniel:************@ep-snowy-lake-123456.us-east-2.aws.neon.tech/holiday_songs?sslmode=require"

import os
import psycopg2
import environ
import requests
import io

# 1. Setup and read the .env file
env = environ.Env()
environ.Env.read_env()

# 2. Establish a connection to the PostgreSQL database
connection_string = env('DATABASE_URL')
conn = psycopg2.connect(connection_string)
cursor = conn.cursor()

# 3. Create the songs table if it doesn't already exist
table_creation_sql = """
CREATE TABLE IF NOT EXISTS public.songs (
    id INT,
    year INT,
    position INT,
    song TEXT,
    artist TEXT,
    chart_date DATE
);
"""
cursor.execute(table_creation_sql)

# 4. Download the holiday_songs.csv file from GitHub
url = "https://github.com/neondatabase/postgres-sample-dbs/raw/main/holiday_songs.csv"
response = requests.get(url)
response.raise_for_status()

# Use StringIO to convert text data into file-like object so it can be read into the database
csv_file = io.StringIO(response.text)

# 5. Load the data from the holiday_songs.csv into the songs table
copy_command = '''
COPY public.songs (id, year, position, song, artist, chart_date)
FROM STDIN WITH (FORMAT CSV, HEADER true, DELIMITER ',');
'''
cursor.copy_expert(copy_command, csv_file)

# 6. Commit any changes and close the connection
conn.commit()
conn.close()

python db.py

import os
import environ
import psycopg2
from urllib.parse import urlparse
from langchain.utilities import SQLDatabase
from langchain.llms import OpenAI
from langchain_experimental.sql import SQLDatabaseChain
from langchain.prompts import PromptTemplate

# Set up and read the .env file
env = environ.Env()
environ.Env.read_env()

# Extract connection details from the .env file
connection_string = env('DATABASE_URL')
parsed_uri = urlparse(connection_string)
username = parsed_uri.username
password = parsed_uri.password
host = parsed_uri.hostname
port = parsed_uri.port or 5432
database = parsed_uri.path[1:]  # remove leading '/'

# Setup database
db = SQLDatabase.from_uri(
    f"postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}"
)

# Setup llm
llm = OpenAI(temperature=0, openai_api_key=os.environ["OPENAI_API_KEY"])

# Define table_info and few_shot_examples
table_info = """public.songs (
    id  integer,
    year integer,
    position integer,
    song text,
    artist text,
    chart_date date
)"""

few_shot_examples = """
- Question: "Who topped the charts in 2018?"
  SQLQuery: SELECT artist, song FROM songs WHERE year = 2018 AND position = 1;

- Question: "Which artist has the most songs that have made it to the top 20 holiday songs list?"
  SQLQuery: SELECT artist, COUNT(*) AS song_count FROM holiday_songs GROUP BY artist ORDER BY song_count DESC LIMIT 1;

- Question: "What is the highest chart position that a holiday song has achieved each year?"
  SQLQuery: SELECT year, MIN(position) AS highest_position FROM holiday_songs GROUP BY year;

- Question: "How many unique songs have charted in the top 20 over the years?"
  SQLQuery: SELECT COUNT(DISTINCT song) FROM holiday_songs;

- Question: "Which song has appeared most frequently in the top 20 holiday songs list?"
  SQLQuery: SELECT song, COUNT(*) AS appearance_count FROM holiday_songs GROUP BY song ORDER BY appearance_count DESC LIMIT 1;

- Question: "In which year did we have the most number of unique artists in the top 20 holiday songs list?"
  SQLQuery: SELECT year, COUNT(DISTINCT artist) AS unique_artists FROM holiday_songs GROUP BY year ORDER BY unique_artists DESC LIMIT 1;
"""

# Define Custom Prompt
TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Use the following format:

Question: "Question here"
SQLQuery: "SQL Query to run"
SQLResult: "Result of the SQLQuery"
Answer: "Final answer here"

Only use the following tables:

{table_info}

Some examples of SQL queries that correspond to questions are:

{few_shot_examples}

Question: {input}"""

CUSTOM_PROMPT = PromptTemplate(
    input_variables=["input", "few_shot_examples", "table_info", "dialect"], template=TEMPLATE
)

# Setup the database chain
db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)

def get_prompt():
    print("Type 'exit' to quit")
    while True:
        prompt = input("Ask a question or type exit to quit: ")

        if prompt.lower() == 'exit':
            print('Exiting...')
            break
        else:
            try:
                question = CUSTOM_PROMPT.format(
                    input=prompt,
                    few_shot_examples=few_shot_examples,
                    table_info=table_info,
                    dialect="PostgreSQL"
                )
                print(db_chain.run(question))
            except Exception as e:
                print(e)

get_prompt()

python app.py