openai-cookbook/examples/vector_databases/neon/neon-postgres-vector-search-pgvector.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vector similarity search using Neon Postgres\n",
    "\n",
    "This notebook guides you through using [Neon Serverless Postgres](https://neon.tech/) as a vector database for OpenAI embeddings. It demonstrates how to:\n",
    "\n",
    "1. Use embeddings created by OpenAI API.\n",
    "2. Store embeddings in a Neon Serverless Postgres database.\n",
    "3. Convert a raw text query to an embedding with OpenAI API.\n",
    "4. Use Neon with the `pgvector` extension to perform vector similarity search."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "Before you begin, ensure that you have the following:\n",
    "\n",
    "1. A Neon Postgres database. You can create an account and set up a project with a ready-to-use `neondb` database in a few simple steps. For instructions, see [Sign up](https://neon.tech/docs/get-started-with-neon/signing-up) and [Create your first project](https://neon.tech/docs/get-started-with-neon/setting-up-a-project).\n",
    "2. A connection string for your Neon 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).\n",
    "3. The `pgvector` extension. Install the extension in Neon by running `CREATE EXTENSION vector;`. For instructions, see [Enable the pgvector extension](https://neon.tech/docs/extensions/pgvector#enable-the-pgvector-extension). \n",
    "4. Your [OpenAI API key](https://platform.openai.com/account/api-keys).\n",
    "5. Python and `pip`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install required modules\n",
    "\n",
    "This notebook requires the `openai`, `psycopg2`, `pandas`, `wget`, and `python-dotenv` packages. You can install them with `pip`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install openai psycopg2 pandas wget python-dotenv"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare your OpenAI API key\n",
    "\n",
    "An OpenAI API key is required to generate vectors for documents and queries.\n",
    "\n",
    "If you do not have an OpenAI API key, obtain one from https://platform.openai.com/account/api-keys.\n",
    "\n",
    "Add the OpenAI API key as an operating system environment variable or provide it for the session when prompted. If you define an environment variable, name the variable `OPENAI_API_KEY`.\n",
    "\n",
    "For information about configuring your OpenAI API key as an environment variable, refer to [Best Practices for API Key Safety](https://help.openai.com/en/articles/5112595-best-practices-for-api-key-safety)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test your OpenAPI key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your OPENAI_API_KEY is ready\n"
     ]
    }
   ],
   "source": [
    "# Test to ensure that your OpenAI API key is defined as an environment variable or provide it when prompted\n",
    "# If you run this notebook locally, you may have to reload the terminal and the notebook to make the environment available\n",
    "\n",
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "# Check if OPENAI_API_KEY is set as an environment variable\n",
    "if os.getenv(\"OPENAI_API_KEY\") is not None:\n",
    "    print(\"Your OPENAI_API_KEY is ready\")\n",
    "else:\n",
    "    # If not, prompt for it\n",
    "    api_key = getpass(\"Enter your OPENAI_API_KEY: \")\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\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to your Neon database\n",
    "\n",
    "Provide your Neon database connection string below or define it in an `.env` file using a `DATABASE_URL` variable. For information about obtaining a Neon connection string, see [Connect from any application](https://neon.tech/docs/connect/connect-from-any-app)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import psycopg2\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load environment variables from .env file\n",
    "load_dotenv()\n",
    "\n",
    "# The connection string can be provided directly here.\n",
    "# Replace the next line with Your Neon connection string.\n",
    "connection_string = \"postgres://<user>:<password>@<hostname>/<dbname>\"\n",
    "\n",
    "# If connection_string is not directly provided above, \n",
    "# then check if DATABASE_URL is set in the environment or .env.\n",
    "if not connection_string:\n",
    "    connection_string = os.environ.get(\"DATABASE_URL\")\n",
    "\n",
    "    # If neither method provides a connection string, raise an error.\n",
    "    if not connection_string:\n",
    "        raise ValueError(\"Please provide a valid connection string either in the code or in the .env file as DATABASE_URL.\")\n",
    "\n",
    "# Connect using the connection string\n",
    "connection = psycopg2.connect(connection_string)\n",
    "\n",
    "# Create a new cursor object\n",
    "cursor = connection.cursor()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test the connection to your database:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your database connection was successful!\n"
     ]
    }
   ],
   "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.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This guide uses pre-computed Wikipedia article embeddings available in the OpenAI Cookbook `examples` directory so that you do not have to compute embeddings with your own OpenAI credits. \n",
    "\n",
    "Import the pre-computed embeddings zip file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'vector_database_wikipedia_articles_embedded.zip'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wget\n",
    "\n",
    "embeddings_url = \"https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip\"\n",
    "\n",
    "# The file is ~700 MB. Importing it will take several minutes.\n",
    "wget.download(embeddings_url)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract the downloaded zip file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The file vector_database_wikipedia_articles_embedded.csv exists in the data directory.\n"
     ]
    }
   ],
   "source": [
    "import zipfile\n",
    "import os\n",
    "import re\n",
    "import tempfile\n",
    "\n",
    "current_directory = os.getcwd()\n",
    "zip_file_path = os.path.join(current_directory, \"vector_database_wikipedia_articles_embedded.zip\")\n",
    "output_directory = os.path.join(current_directory, \"../../data\")\n",
    "\n",
    "with zipfile.ZipFile(zip_file_path, \"r\") as zip_ref:\n",
    "    zip_ref.extractall(output_directory)\n",
    "\n",
    "\n",
    "# Check to see if the csv file was extracted\n",
    "file_name = \"vector_database_wikipedia_articles_embedded.csv\"\n",
    "data_directory = os.path.join(current_directory, \"../../data\")\n",
    "file_path = os.path.join(data_directory, file_name)\n",
    "\n",
    "\n",
    "if os.path.exists(file_path):\n",
    "    print(f\"The csv file {file_name} exists in the data directory.\")\n",
    "else:\n",
    "    print(f\"The csv file {file_name} does not exist in the data directory.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a table and add indexes for your vector embeddings\n",
    "\n",
    "The vector table created in your database is called **articles**. Each object has **title** and **content** vectors. \n",
    "\n",
    "An index is defined on both the **title** and **content** vector columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_table_sql = '''\n",
    "CREATE TABLE IF NOT EXISTS public.articles (\n",
    "    id INTEGER NOT NULL,\n",
    "    url TEXT,\n",
    "    title TEXT,\n",
    "    content TEXT,\n",
    "    title_vector vector(1536),\n",
    "    content_vector vector(1536),\n",
    "    vector_id INTEGER\n",
    ");\n",
    "\n",
    "ALTER TABLE public.articles ADD PRIMARY KEY (id);\n",
    "'''\n",
    "\n",
    "# SQL statement for creating indexes\n",
    "create_indexes_sql = '''\n",
    "CREATE INDEX ON public.articles USING ivfflat (content_vector) WITH (lists = 1000);\n",
    "\n",
    "CREATE INDEX ON public.articles USING ivfflat (title_vector) WITH (lists = 1000);\n",
    "'''\n",
    "\n",
    "# Execute the SQL statements\n",
    "cursor.execute(create_table_sql)\n",
    "cursor.execute(create_indexes_sql)\n",
    "\n",
    "# Commit the changes\n",
    "connection.commit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data\n",
    "\n",
    "Load the pre-computed vector data into your `articles` table from the `.csv` file. There are 25000 records, so expect the operation to take several minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "# Path to your local CSV file\n",
    "csv_file_path = '../../data/vector_database_wikipedia_articles_embedded.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.articles (id, url, title, content, title_vector, content_vector, vector_id)\n",
    "FROM STDIN WITH (FORMAT CSV, HEADER true, DELIMITER ',');\n",
    "'''\n",
    "\n",
    "# Execute the COPY command using copy_expert\n",
    "cursor.copy_expert(copy_command, modified_lines)\n",
    "\n",
    "# Commit the changes\n",
    "connection.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the number of records to ensure the data has been been loaded. There should be 25000 records."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count:25000\n"
     ]
    }
   ],
   "source": [
    "# Check the size of the data\n",
    "count_sql = \"\"\"select count(*) from public.articles;\"\"\"\n",
    "cursor.execute(count_sql)\n",
    "result = cursor.fetchone()\n",
    "print(f\"Count:{result[0]}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Search your data\n",
    "\n",
    "After the data is stored in your Neon database, you can query the data for nearest neighbors. \n",
    "\n",
    "Start by defining the `query_neon` function, which is executed when you run the vector similarity search. The function creates an embedding based on the user's query, prepares the SQL query, and runs the SQL query with the embedding. The pre-computed embeddings that you loaded into your database were created with `text-embedding-3-small` OpenAI model, so you must use the same model to create an embedding for the similarity search.\n",
    "\n",
    "A `vector_name` parameter is provided that allows you to search based on \"title\" or \"content\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_neon(query, collection_name, vector_name=\"title_vector\", top_k=20):\n",
    "\n",
    "    # Create an embedding vector from the user query\n",
    "    embedded_query = openai.Embedding.create(\n",
    "        input=query,\n",
    "        model=\"text-embedding-3-small\",\n",
    "    )[\"data\"][0][\"embedding\"]\n",
    "\n",
    "    # Convert the embedded_query to PostgreSQL compatible format\n",
    "    embedded_query_pg = \"[\" + \",\".join(map(str, embedded_query)) + \"]\"\n",
    "\n",
    "    # Create the SQL query\n",
    "    query_sql = f\"\"\"\n",
    "    SELECT id, url, title, l2_distance({vector_name},'{embedded_query_pg}'::VECTOR(1536)) AS similarity\n",
    "    FROM {collection_name}\n",
    "    ORDER BY {vector_name} <-> '{embedded_query_pg}'::VECTOR(1536)\n",
    "    LIMIT {top_k};\n",
    "    \"\"\"\n",
    "    # Execute the query\n",
    "    cursor.execute(query_sql)\n",
    "    results = cursor.fetchall()\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run a similarity search based on `title_vector` embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. Greek mythology (Score: 0.998)\n",
      "2. Roman mythology (Score: 0.7)\n",
      "3. Greek underworld (Score: 0.637)\n",
      "4. Mythology (Score: 0.635)\n",
      "5. Classical mythology (Score: 0.629)\n",
      "6. Japanese mythology (Score: 0.615)\n",
      "7. Norse mythology (Score: 0.569)\n",
      "8. Greek language (Score: 0.566)\n",
      "9. Zeus (Score: 0.534)\n",
      "10. List of mythologies (Score: 0.531)\n",
      "11. Jupiter (mythology) (Score: 0.53)\n",
      "12. Greek (Score: 0.53)\n",
      "13. Gaia (mythology) (Score: 0.526)\n",
      "14. Titan (mythology) (Score: 0.522)\n",
      "15. Mercury (mythology) (Score: 0.521)\n",
      "16. Ancient Greece (Score: 0.52)\n",
      "17. Greek alphabet (Score: 0.52)\n",
      "18. Venus (mythology) (Score: 0.515)\n",
      "19. Pluto (mythology) (Score: 0.515)\n",
      "20. Athena (Score: 0.514)\n"
     ]
    }
   ],
   "source": [
    "# Query based on `title_vector` embeddings\n",
    "import openai\n",
    "\n",
    "query_results = query_neon(\"Greek mythology\", \"Articles\")\n",
    "for i, result in enumerate(query_results):\n",
    "    print(f\"{i + 1}. {result[2]} (Score: {round(1 - result[3], 3)})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run a similarity search based on `content_vector` embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. 222 BC (Score: 0.489)\n",
      "2. Trojan War (Score: 0.458)\n",
      "3. Peloponnesian War (Score: 0.456)\n",
      "4. History of the Peloponnesian War (Score: 0.449)\n",
      "5. 430 BC (Score: 0.441)\n",
      "6. 168 BC (Score: 0.436)\n",
      "7. Ancient Greece (Score: 0.429)\n",
      "8. Classical Athens (Score: 0.428)\n",
      "9. 499 BC (Score: 0.427)\n",
      "10. Leonidas I (Score: 0.426)\n",
      "11. Battle (Score: 0.421)\n",
      "12. Greek War of Independence (Score: 0.421)\n",
      "13. Menelaus (Score: 0.419)\n",
      "14. Thebes, Greece (Score: 0.417)\n",
      "15. Patroclus (Score: 0.417)\n",
      "16. 427 BC (Score: 0.416)\n",
      "17. 429 BC (Score: 0.413)\n",
      "18. August 2 (Score: 0.412)\n",
      "19. Ionia (Score: 0.411)\n",
      "20. 323 (Score: 0.409)\n"
     ]
    }
   ],
   "source": [
    "# Query based on `content_vector` embeddings\n",
    "query_results = query_neon(\"Famous battles in Greek history\", \"Articles\", \"content_vector\")\n",
    "for i, result in enumerate(query_results):\n",
    "    print(f\"{i + 1}. {result[2]} (Score: {round(1 - result[3], 3)})\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}