Semantic search using Supabase Vector

Dec 4, 2023
Open in Github

The purpose of this guide is to demonstrate how to store OpenAI embeddings in Supabase Vector (Postgres + pgvector) for the purposes of semantic search.

Supabase is an open-source Firebase alternative built on top of Postgres, a production-grade SQL database. Since Supabase Vector is built on pgvector, you can store your embeddings within the same database that holds the rest of your application data. When combined with pgvector's indexing algorithms, vector search remains fast at large scales.

Supabase adds an ecosystem of services and tools to make app development as quick as possible (such as an auto-generated REST API). We'll use these services to store and query embeddings within Postgres.

This guide covers:

  1. Setting up your database
  2. Creating a SQL table that can store vector data
  3. Generating OpenAI embeddings using OpenAI's JavaScript client
  4. Storing the embeddings in your SQL table using the Supabase JavaScript client
  5. Performing semantic search over the embeddings using a Postgres function and the Supabase JavaScript client

Setup database

First head over to https://database.new to provision your Supabase database. This will create a Postgres database on the Supabase cloud platform. Alternatively, you can follow the local development options if you prefer to run your database locally using Docker.

In the studio, jump to the SQL editor and execute the following SQL to enable pgvector:

-- Enable the pgvector extension
create extension if not exists vector;

In a production application, the best practice is to use database migrations so that all SQL operations are managed within source control. To keep things simple in this guide, we'll execute queries directly in the SQL Editor. If you are building a production app, feel free to move these into a database migration.

Create a vector table

Next we'll create a table to store documents and embeddings. In the SQL Editor, run:

create table documents (
  id bigint primary key generated always as identity,
  content text not null,
  embedding vector (1536) not null
);

Since Supabase is built on Postgres, we're just using regular SQL here. You can modify this table however you like to better fit your application. If you have existing database tables, you can simply add a new vector column to the appropriate table.

The important piece to understand is the vector data type, which is a new data type that became available when we enabled the pgvector extension earlier. The size of the vector (1536 here) represents the number of dimensions in the embedding. Since we're using OpenAI's text-embedding-3-small model in this example, we set the vector size to 1536.

Let's go ahead and create a vector index on this table so that future queries remain performant as the table grows:

create index on documents using hnsw (embedding vector_ip_ops);

This index uses the HNSW algorithm to index vectors stored in the embedding column, and specifically when using the inner product operator (<#>). We'll explain more about this operator later when we implement our match function.

Let's also follow security best practices by enabling row level security on the table:

alter table documents enable row level security;

This will prevent unauthorized access to this table through the auto-generated REST API (more on this shortly).

Generate OpenAI embeddings

This guide uses JavaScript to generate embeddings, but you can easily modify it to use any language supported by OpenAI.

If you are using JavaScript, feel free to use whichever server-side JavaScript runtime that you prefer (Node.js, Deno, Supabase Edge Functions).

If you're using Node.js, first install openai as a dependency:

npm install openai

then import it:

import OpenAI from "openai";

If you're using Deno or Supabase Edge Functions, you can import openai directly from a URL:

import OpenAI from "https://esm.sh/openai@4";

In this example we import from https://esm.sh which is a CDN that automatically fetches the respective NPM module for you and serves it over HTTP.

Next we'll generate an OpenAI embedding using text-embedding-3-small:

const openai = new OpenAI();
 
const input = "The cat chases the mouse";
 
const result = await openai.embeddings.create({
  input,
  model: "text-embedding-3-small",
});
 
const [{ embedding }] = result.data;

Remember that you will need an OpenAI API key to interact with the OpenAI API. You can pass this as an environment variable called OPENAI_API_KEY, or manually set it when you instantiate your OpenAI client:

const openai = new OpenAI({
  apiKey: "<openai-api-key>",
});

Remember: Never hard-code API keys in your code. Best practice is to either store it in a .env file and load it using a library like dotenv or load it from an external key management system.

Store embeddings in database

Supabase comes with an auto-generated REST API that dynamically builds REST endpoints for each of your tables. This means you don't need to establish a direct Postgres connection to your database - instead you can interact with it simply using by the REST API. This is especially useful in serverless environments that run short-lived processes where re-establishing a database connection every time can be expensive.

Supabase comes with a number of client libraries to simplify interaction with the REST API. In this guide we'll use the JavaScript client library, but feel free to adjust this to your preferred language.

If you're using Node.js, install @supabase/supabase-js as a dependency:

npm install @supabase/supabase-js

then import it:

import { createClient } from "@supabase/supabase-js";

If you're using Deno or Supabase Edge Functions, you can import @supabase/supabase-js directly from a URL:

import { createClient } from "https://esm.sh/@supabase/supabase-js@2";

Next we'll instantiate our Supabase client and configure it so that it points to your Supabase project. In this guide we'll store a reference to your Supabase URL and key in a .env file, but feel free to modify this based on how your application handles configuration.

If you are using Node.js or Deno, add your Supabase URL and service role key to a .env file. If you are using the cloud platform, you can find these from your Supabase dashboard settings page. If you're running Supabase locally, you can find these by running npx supabase status in a terminal.

.env

SUPABASE_URL=<supabase-url>
SUPABASE_SERVICE_ROLE_KEY=<supabase-service-role-key>

If you are using Supabase Edge Functions, these environment variables are automatically injected into your function for you so you can skip the above step.

Next we'll pull these environment variables into our app.

In Node.js, install the dotenv dependency:

npm install dotenv

And retrieve the environment variables from process.env:

import { config } from "dotenv";
 
// Load .env file
config();
 
const supabaseUrl = process.env["SUPABASE_URL"];
const supabaseServiceRoleKey = process.env["SUPABASE_SERVICE_ROLE_KEY"];

In Deno, load the .env file using the dotenv standard library:

import { load } from "https://deno.land/std@0.208.0/dotenv/mod.ts";
 
// Load .env file
const env = await load();
 
const supabaseUrl = env["SUPABASE_URL"];
const supabaseServiceRoleKey = env["SUPABASE_SERVICE_ROLE_KEY"];

In Supabase Edge Functions, simply load the injected environment variables directly:

const supabaseUrl = Deno.env.get("SUPABASE_URL");
const supabaseServiceRoleKey = Deno.env.get("SUPABASE_SERVICE_ROLE_KEY");

Next let's instantiate our supabase client:

const supabase = createClient(supabaseUrl, supabaseServiceRoleKey, {
  auth: { persistSession: false },
});

From here we use the supabase client to insert our text and embedding (generated earlier) into the database:

const { error } = await supabase.from("documents").insert({
  content: input,
  embedding,
});

In production, best practice would be to check the response error to see if there were any problems inserting the data and handle it accordingly.

Finally let's perform semantic search over the embeddings in our database. At this point we'll assume your documents table has been filled with multiple records that we can search over.

Let's create a match function in Postgres that performs the semantic search query. Execute the following in the SQL Editor:

create function match_documents (
  query_embedding vector (1536),
  match_threshold float,
)
returns setof documents
language plpgsql
as $$
begin
  return query
  select *
  from documents
  where documents.embedding <#> query_embedding < -match_threshold
  order by documents.embedding <#> query_embedding;
end;
$$;

This function accepts a query_embedding which represents the embedding generated from the search query text (more on this shortly). It also accepts a match_threshold which specifies how similar the document embeddings have to be in order for query_embedding to count as a match.

Inside the function we implement the query which does two things:

  • Filters the documents to only include those who's embeddings match within the above match_threshold. Since the <#> operator performs the negative inner product (versus positive inner product), we negate the similarity threshold before comparing. This means a match_threshold of 1 is most similar, and -1 is most dissimilar.
  • Orders the documents by negative inner product (<#>) ascending. This allows us to retrieve documents that match closest first.

Since OpenAI embeddings are normalized, we opted to use inner product (<#>) because it is slightly more performant than other operators like cosine distance (<=>). It is important to note though this only works because the embeddings are normalized - if they weren't, cosine distance should be used.

Now we can call this function from our application using the supabase.rpc() method:

const query = "What does the cat chase?";
 
// First create an embedding on the query itself
const result = await openai.embeddings.create({
  input: query,
  model: "text-embedding-3-small",
});
 
const [{ embedding }] = result.data;
 
// Then use this embedding to search for matches
const { data: documents, error: matchError } = await supabase
  .rpc("match_documents", {
    query_embedding: embedding,
    match_threshold: 0.8,
  })
  .select("content")
  .limit(5);

In this example, we set a match threshold to 0.8. Adjust this threshold based on what works best with your data.

Note that since match_documents returns a set of documents, we can treat this rpc() like a regular table query. Specifically this means we can chain additional commands to this query, like select() and limit(). Here we select just the columns we care about from the documents table (content), and we limit the number of documents returned (max 5 in this example).

At this point you have a list of documents that matched the query based on semantic relationship, ordered by most similar first.

Next steps

You can use this example as the foundation for other semantic search techniques, like retrieval augmented generation (RAG).

For more information on OpenAI embeddings, read the Embedding docs.

For more information on Supabase Vector, read the AI & Vector docs.