Perform Vector Similarity Search with Neon Postgres

Vector similarity search using Neon Postgres

This notebook guides you through using Neon Serverless Postgres as a vector database for OpenAI embeddings. It demonstrates how to:

1. Use embeddings created by OpenAI API.
2. Store embeddings in a Neon Serverless Postgres database.
3. Convert a raw text query to an embedding with OpenAI API.
4. Use Neon with the pgvector extension to perform vector similarity search.

Prerequisites

Before you begin, ensure that you have the following:

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 and Create your first project.
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.
3. The pgvector extension. Install the extension in Neon by running CREATE EXTENSION vector;. For instructions, see Enable the pgvector extension.
4. Your OpenAI API key.
5. Python and pip.

Install required modules

This notebook requires the openai, psycopg2, pandas, wget, and python-dotenv packages. You can install them with pip:

Prepare your OpenAI API key

An OpenAI API key is required to generate vectors for documents and queries.

If you do not have an OpenAI API key, obtain one from https://platform.openai.com/account/api-keys.

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.

For information about configuring your OpenAI API key as an environment variable, refer to Best Practices for API Key Safety.

Test your OpenAPI key

Connect to your Neon database

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.

Test the connection to your database:

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.

Import the pre-computed embeddings zip file:

Extract the downloaded zip file:

Create a table and add indexes for your vector embeddings

The vector table created in your database is called articles. Each object has title and content vectors.

An index is defined on both the title and content vector columns.

Load the data

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.

Check the number of records to ensure the data has been been loaded. There should be 25000 records.

Search your data

After the data is stored in your Neon database, you can query the data for nearest neighbors.

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.

A vector_name parameter is provided that allows you to search based on "title" or "content".

Run a similarity search based on title_vector embeddings:

Run a similarity search based on content_vector embeddings: