chat_with_your_database_using_langchain.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "**Prerequisites**\n",
        "\n",
        "Before you begin, ensure that you have the following:\n",
        "\n",
        "- A database named `titanic`. You can create one in just a few clicks using the Neon console. See [Create a database](https://neon.tech/docs/manage/databases#create-a-database).\n",
        "- A connection string for your `titanic` database. You can copy it from the **Connection Details** widget on the Neon Dashboard. See [Connect from any application](https://neon.tech/docs/connect/connect-from-any-app) for instructions.\n",
        "- Your OpenAI API key. If you do not have an OpenAI API key, obtain one from [https://platform.openai.com/account/api-keys](https://platform.openai.com/account/api-keys)."
      ],
      "metadata": {
        "id": "gMUJVW5scPk0"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Install the required modules**\n",
        "\n",
        "This notebook requires the `LangChain`, `LangChain-experimental`, and `openai` modules. Run this code block to install them."
      ],
      "metadata": {
        "id": "9850-ana3z1S"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install LangChain LangChain-experimental openai"
      ],
      "metadata": {
        "id": "176iwLWL4IjN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Provide your OpenAI API key**\n",
        "\n",
        "Run this code block and provide your OpenAI API key when prompted."
      ],
      "metadata": {
        "id": "bhtUoDQW4eEd"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xA7X9zpGkXyU",
        "outputId": "fae52a47-3f6f-4fc6-8495-bc95669d92c1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Enter your OPENAI_API_KEY: ··········\n",
            "Your OPENAI_API_KEY is now available for this session\n"
          ]
        }
      ],
      "source": [
        "from getpass import getpass\n",
        "import os\n",
        "\n",
        "# Directly prompt for the API key\n",
        "api_key = getpass(\"Enter your OPENAI_API_KEY: \")\n",
        "\n",
        "if api_key:\n",
        "    print(\"Your OPENAI_API_KEY is now available for this session\")\n",
        "    # Optionally, you can set it as an environment variable for the current session\n",
        "    os.environ[\"OPENAI_API_KEY\"] = api_key\n",
        "else:\n",
        "    print(\"You did not enter your OPENAI_API_KEY\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Provide your database connection string and open a cursor**\n",
        "\n",
        "Input the connection string for your Neon database and run the code block."
      ],
      "metadata": {
        "id": "RNE0nnog5HHC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import psycopg2\n",
        "from urllib.parse import urlparse\n",
        "\n",
        "# Provide your Neon connection string\n",
        "connection_string = \"postgres://<user>:<password>@<hostname>/<dbname>\"\n",
        "\n",
        "# Extract details from connection string\n",
        "parsed_uri = urlparse(connection_string)\n",
        "username = parsed_uri.username\n",
        "password = parsed_uri.password\n",
        "host = parsed_uri.hostname\n",
        "port = parsed_uri.port or 5432\n",
        "database = parsed_uri.path[1:]  # remove leading '/'\n",
        "\n",
        "# Connect using the connection string\n",
        "connection = psycopg2.connect(connection_string)\n",
        "\n",
        "# Create a new cursor object\n",
        "cursor = connection.cursor()"
      ],
      "metadata": {
        "id": "ZAsqq2KPrcxp"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Test your database connection**\n",
        "\n",
        "Run this code block to test your database connection."
      ],
      "metadata": {
        "id": "qhU1omej5SsL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Execute this query to test the database connection\n",
        "cursor.execute(\"SELECT 1;\")\n",
        "result = cursor.fetchone()\n",
        "\n",
        "# Check the query result\n",
        "if result == (1,):\n",
        "    print(\"Your database connection was successful!\")\n",
        "else:\n",
        "    print(\"Your connection failed.\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kn6wKSzZr0rV",
        "outputId": "eb47e5a5-2512-4159-feaf-9e4d0d0e139a"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your database connection was successful!\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Create a table for passenger data**\n",
        "\n",
        "Run this code block to create a `passenger` table in your `titanic` database."
      ],
      "metadata": {
        "id": "82D68vQ95Wml"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "create_table_sql = '''\n",
        "CREATE TABLE public.passenger (\n",
        "    passengerid integer NOT NULL,\n",
        "    survived double precision,\n",
        "    pclass integer,\n",
        "    name text,\n",
        "    sex text,\n",
        "    age double precision,\n",
        "    sibsp integer,\n",
        "    parch integer,\n",
        "    ticket text,\n",
        "    fare double precision,\n",
        "    cabin text,\n",
        "    embarked text,\n",
        "    wikiid double precision,\n",
        "    name_wiki text,\n",
        "    age_wiki double precision,\n",
        "    hometown text,\n",
        "    boarded text,\n",
        "    destination text,\n",
        "    lifeboat text,\n",
        "    body text,\n",
        "    class integer\n",
        ");\n",
        "'''\n",
        "\n",
        "# Execute the SQL statement\n",
        "cursor.execute(create_table_sql)\n",
        "\n",
        "# Commit the changes\n",
        "connection.commit()"
      ],
      "metadata": {
        "id": "LT_H3_SLr4LY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Load data into the passenger table**\n",
        "\n",
        "Run this code block to load data into the `passenger` table in your `titanic` database from a `titanic.csv` file hosted on GitHub."
      ],
      "metadata": {
        "id": "vSanI0Jn5eGg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import io\n",
        "\n",
        "# Download the file from GitHub\n",
        "!wget https://github.com/neondatabase/postgres-sample-dbs/raw/main/titanic.csv -O titanic.csv\n",
        "\n",
        "# Path to the downloaded CSV file\n",
        "csv_file_path = 'titanic.csv'\n",
        "\n",
        "# Define a generator function to process the csv file\n",
        "def process_file(file_path):\n",
        "    with open(file_path, 'r', encoding='utf-8') as file:\n",
        "        for line in file:\n",
        "            yield line\n",
        "\n",
        "# Create a StringIO object to store the modified lines\n",
        "modified_lines = io.StringIO(''.join(list(process_file(csv_file_path))))\n",
        "\n",
        "# Create the COPY command for copy_expert\n",
        "copy_command = '''\n",
        "COPY public.passenger (passengerid, survived, pclass, name, sex, age, sibsp, parch, ticket, fare, cabin, embarked, wikiid, name_wiki, age_wiki, hometown, boarded, destination, lifeboat, body, class)\n",
        "FROM STDIN WITH (FORMAT CSV, HEADER true, DELIMITER ',');\n",
        "'''\n",
        "\n",
        "# Assuming you have already set up a connection and cursor...\n",
        "cursor.copy_expert(copy_command, modified_lines)\n",
        "\n",
        "# Commit the changes\n",
        "connection.commit()"
      ],
      "metadata": {
        "id": "nAKxX8m2z8h-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Prompt for a question**\n",
        "\n",
        "Run this code block to ask your database a question. When prompted, enter a question like, \"How many passengers survived?\" or \"What was the average age of passengers?\""
      ],
      "metadata": {
        "id": "iN4IlPm73wWb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.utilities import SQLDatabase\n",
        "from langchain.llms import OpenAI\n",
        "from langchain_experimental.sql import SQLDatabaseChain\n",
        "from langchain.prompts import PromptTemplate\n",
        "\n",
        "# Setup database\n",
        "db = SQLDatabase.from_uri(\n",
        "    f\"postgresql+psycopg2://{username}:{password}@{host}:{port}/{database}\"\n",
        ")\n",
        "\n",
        "# Setup llm\n",
        "llm = OpenAI(temperature=0, openai_api_key=os.environ[\"OPENAI_API_KEY\"])\n",
        "\n",
        "# Define table_info and few_shot_examples\n",
        "table_info = \"\"\"public.passenger (\n",
        "    passengerid integer NOT NULL,\n",
        "    survived double precision,\n",
        "    pclass integer,\n",
        "    name text,\n",
        "    sex text,\n",
        "    age double precision,\n",
        "    sibsp integer,\n",
        "    parch integer,\n",
        "    ticket text,\n",
        "    fare double precision,\n",
        "    cabin text,\n",
        "    embarked text,\n",
        "    wikiid double precision,\n",
        "    name_wiki text,\n",
        "    age_wiki double precision,\n",
        "    hometown text,\n",
        "    boarded text,\n",
        "    destination text,\n",
        "    lifeboat text,\n",
        "    body text,\n",
        "    class integer\n",
        ")\"\"\"\n",
        "\n",
        "few_shot_examples = \"\"\"\n",
        "- Question: \"How many passengers survived?\"\n",
        "  SQLQuery: \"SELECT COUNT(*) FROM public.passenger WHERE survived = 1;\"\n",
        "\n",
        "- Question: \"What was the average age of passengers?\"\n",
        "  SQLQuery: \"SELECT AVG(age) FROM public.passenger;\"\n",
        "\n",
        "- Question: \"How many male and female passengers were there?\"\n",
        "  SQLQuery: \"SELECT sex, COUNT(*) FROM public.passenger GROUP BY sex;\"\n",
        "\n",
        "- Question: \"Which passenger had the highest fare?\"\n",
        "  SQLQuery: \"SELECT name, fare FROM public.passenger WHERE fare IS NOT NULL ORDER BY fare DESC LIMIT 1;\"\n",
        "\n",
        "- Question: \"How many passengers boarded from each location?\"\n",
        "  SQLQuery: \"SELECT embarked, COUNT(*) FROM public.passenger GROUP BY embarked;\"\n",
        "\n",
        "- Question: \"Who is the oldest passenger and what was their age?\":\n",
        "  SQLQuery: SELECT name, age FROM public.passenger WHERE age IS NOT NULL ORDER BY age DESC LIMIT 1;\n",
        "\"\"\"\n",
        "\n",
        "# Define Custom Prompt\n",
        "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.\n",
        "Use the following format:\n",
        "\n",
        "Question: \"Question here\"\n",
        "SQLQuery: \"SQL Query to run\"\n",
        "SQLResult: \"Result of the SQLQuery\"\n",
        "Answer: \"Final answer here\"\n",
        "\n",
        "Only use the following tables:\n",
        "\n",
        "{table_info}\n",
        "\n",
        "Some examples of SQL queries that correspond to questions are:\n",
        "\n",
        "{few_shot_examples}\n",
        "\n",
        "Question: {input}\"\"\"\n",
        "\n",
        "CUSTOM_PROMPT = PromptTemplate(\n",
        "    input_variables=[\"input\", \"few_shot_examples\", \"table_info\", \"dialect\"], template=TEMPLATE\n",
        ")\n",
        "\n",
        "# Setup the database chain\n",
        "# Setup the database chain\n",
        "db_chain = SQLDatabaseChain(llm=llm, database=db, verbose=True)\n",
        "\n",
        "def get_prompt():\n",
        "    print(\"Type 'exit' to quit\")\n",
        "    while True:\n",
        "        prompt = input(\"Ask a question or type exit to quit: \")\n",
        "\n",
        "        if prompt.lower() == 'exit':\n",
        "            print('Exiting...')\n",
        "            break\n",
        "        else:\n",
        "            try:\n",
        "                question = CUSTOM_PROMPT.format(\n",
        "                    input=prompt,\n",
        "                    few_shot_examples=few_shot_examples,\n",
        "                    table_info=table_info,\n",
        "                    dialect=\"PostgreSQL\"\n",
        "                )\n",
        "                print(db_chain.run(question))\n",
        "            except Exception as e:\n",
        "                print(e)\n",
        "\n",
        "get_prompt()"
      ],
      "metadata": {
        "id": "F0ZdmiGFAAsr"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}