Postgres with pg_vector extension, Docker setup

[ndrean]

Install Vector extension to Postgres with Docker

To add the pg_vector extension to your PostgreSQL container, you'll need to use a PostgreSQL image that includes this extension.

The official PostgreSQL image doesn't include pg_vector by default, so we'll use a custom image that has it pre-installed.

Create a Dockerfile with the following content:

FROM postgres:16

RUN apt-get update && apt-get install -y \
    git \
    build-essential \
    postgresql-server-dev-16

RUN git clone https://github.com/pgvector/pgvector.git && \
    cd pgvector && \
    make && \
    make install

CMD ["postgres"]

Build the custom image named "postgres-with-vector":

docker build -t postgres-with-vector .

Run a container named "postgres-rag" with this custom "postgres-with-vector" image, create the database "rag_example", and open the port 5432 for a Postgrex connection:

docker run \
  -it --rm \
  --name postgres-rag \
  -e POSTGRES_PASSWORD=secret \
  -e POSTGRES_DB=rag_example \
  -p 5432:5432 \
postgres-with-vector

In another terminal, connect to the running "postgres-rag" container and execute psql to connect to the "rag_example" database:

docker exec -it postgres-rag psql -U postgres -d rag_example

With Elixir, we use the Ecto.Repo behaviour. This give an easy DSL (instead of our SQL vis Postgrex):

Postgrex.Types.define(RAG.PostgrexTypes, Pgvector.extensions() ++ Ecto.Adapters.Postgres.extensions(), [])

defmodule RAG.Repo do
  use Ecto.Repo,
    otp_app: :rag,
    adapter: Ecto.Adapters.Postgres
end

defmodule RAG.Document do
  use Ecto.Schema

  schema "documents" do
    field :content, :string
    field :embedding, Pgvector.Ecto.Vector
  end
end

{:ok, pg} = 
  RAG.Repo.start_link(
    hostname: "localhost",
    username: "postgres",
    password: "secret",
    database: "rag_example",
    types: RAG.PostgrexTypes
  )

Note that we can use directly Postgrex.query! commands instead of using the module RAG.Repo. If we use Postgrex, we can do something like:

{:ok, pg} = Postgrex.start_link(database: "rag_example", type: "RAG.PostgrexTypes")

Postgrex.query!(pg, "create extension if not exists vector;", [])
Postgrex.query!(pg, "drop table if exists documents;", [])
Postgrex.query!(pg, "create table documents ....", [])

We use our "Repo" module".

RAG.Repo.query!("create extension if not exists vector;")

RAG.Repo.query!("drop table if exists documents;")

We check in the terminal that the index HNSW method is available:

rag_example=# select * from pg_am where amname='hnsw'

16450 | hnsw   | hnswhandler | i

We create a table with two columns, "content" and "embedding" where the datatypes are respectively "text" and "vector(384)". The later is because we will be using an embedding model with 384 dimensions (see further).

We create an hnsw index on the "embedding" column using the "cosine" distance.

cf documentation: an HNSW index creates a multilayer graph. It has better query performance than IVFFlat (in terms of speed-recall tradeoff), but has slower build times and uses more memory. Also, an index can be created without any data in the table

RAG.Repo.query!("""
  CREATE TABLE IF NOT EXISTS documents (
    id SERIAL PRIMARY KEY,
    content TEXT,
    embedding vector(384)
  )
""")

RAG.Repo.query!(
  "create index if not exists embedding_idx on documents using hnsw (embedding vector_cosine_ops);"
)

Check in the terminal (that runs psql in the container) the details of the created table "documents" and the indexes we created:

rag_example=# \d documents

 id        | integer     |   | not null | nextval('documents_id_seq'::regclass)
 content   | text        |   |          |
 embedding | vector(384) |   |          |

rag_example=# select * from pg_indexes where tablename='documents'

 public | documents | documents_pkey           |  | CREATE UNIQUE INDEX documents_pkey ON public.documents USING btree (id)
 public | documents | documents_embedding_idx  |  | CREATE INDEX documents_embedding_idx ON public.documents USING hnsw (embedding vector_cosine_ops)