This is the multi-page printable view of this section. .

Return to the regular view of this page.

Integration

Integrate pgvecto.rs with LangChain, LlamaIndex, and Citus distributed PostgreSQL.

1: LangChain
2: LlamaIndex
3: Citus

1 - LangChain

This document describes the LangChain integration.

LangChain is a framework for developing applications powered by language models. It is a set of tools that allows you to fine-tune, deploy, and use language models in production.

pgvecto.rs provides a LangChain integration that allows you to retrieve the most similar vectors in LangChain.

Install dependencies

Some dependencies are required to use the LangChain integration:

pip install langchain langchain-openai pgvecto-rs

You could start the postgres instance with pgvecto.rs extension in a docker container:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.0

Then you can connect to the database using the psql command line tool. The default username is postgres, and the default password is mysecretpassword.

psql postgresql://postgres:mysecretpassword@localhost:5432/postgres

Run the following SQL to ensure the extension is enabled.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

Create the database and load documents

We will show how to use pgvecto.rs in LangChain to retrieve the most similar vectors.

First, you need to create the text loader and the text splitter to split the text into chunks. We use the markdown file getting-started/overview.md as an example.

## Loading Environment Variables
from dotenv import load_dotenv

load_dotenv()

from langchain.docstore.document import Document
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores.pgvecto_rs import PGVecto_rs
from langchain_openai import OpenAIEmbeddings


loader = TextLoader("./src/getting-started/overview.md")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings()

Then, we will create the PGVecto_rs instance and load the documents into the database.

## PGVecto.rs needs the connection string to the database.
## We will load it from the environment variables.
import os

PORT = os.getenv("DB_PORT", 5432)
HOST = os.getenv("DB_HOST", "localhost")
USER = os.getenv("DB_USER", "postgres")
PASS = os.getenv("DB_PASS", "mysecretpassword")
DB_NAME = os.getenv("DB_NAME", "postgres")

# Run tests with shell:
URL = "postgresql+psycopg://{username}:{password}@{host}:{port}/{db_name}".format(
    port=PORT,
    host=HOST,
    username=USER,
    password=PASS,
    db_name=DB_NAME,
)

# The pgvectors Module will try to create a table with the name of the collection.
# So, make sure that the collection name is unique and the user has the permission to create a table.

COLLECTION_NAME = "state_of_the_union_test"

db = PGVecto_rs.from_documents(
    embedding=embeddings,
    documents=docs,
    collection_name=COLLECTION_NAME,
    db_url=URL,
)

Query index

Finally, we can retrieve the most similar chunks in LangChain.

query = "What is pgvecto.rs"
docs_with_score = db.similarity_search_with_score(query)

for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)

The output is:

Created a chunk of size 1181, which is longer than the specified 1000
--------------------------------------------------------------------------------
Score:  0.25059962
# Overview

An introduction to the pgvecto.rs.

## What is pgvecto.rs

pgvecto.rs is a Postgres extension that provides vector similarity search functions. It is written in Rust and based on [pgrx](https://github.com/tcdi/pgrx). It is currently in the beta status, we invite you to try it out in production and provide us with feedback. Read more at [📝our launch blog](https://modelz.ai/blog/pgvecto-rs).

## Why use pgvecto.rs
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
Score:  0.29536954
- 💃 **Easy to use**: pgvecto.rs is a Postgres extension, which means that you can use it directly within your existing database. This makes it easy to integrate into your existing workflows and applications.
- 🔗 **Async indexing**: pgvecto.rs's index is asynchronously constructed by the background threads and does not block insertions and always ready for new queries.
- 🥅 **Filtering**: pgvecto.rs supports filtering. You can set conditions when searching or retrieving points. This is the missing feature of other postgres extensions.
- 🧮 **Quantization**: pgvecto.rs supports scalar quantization and product quantization up to 64x.
- 🦀 **Rewrite in Rust**: Rust's strict compile-time checks ensure memory safety, reducing the risk of bugs and security issues commonly associated with C extensions.

## Comparison with pgvector
--------------------------------------------------------------------------------
--------------------------------------------------------------------------------
Score:  0.35845917
More details at [📝`pgvecto.rs` vs. pgvector](../faqs/comparison-pgvector.md).

## Quick start

For new users, we recommend using the [Docker image](https://hub.docker.com/r/tensorchord/pgvecto-rs) to get started quickly.
...

Initialize from an existing database

Above, we created a vector store from scratch. However, often times we want to work with an existing vector store. In order to do that, we can initialize it directly.

db = PGVecto_rs(
    embedding=embeddings,
    collection_name=COLLECTION_NAME,
    # OpenAI embedding has 1536 dimensions.
    dimension=1536,
    db_url=URL,
)

Then, we can add vectors to the store and query:

db.add_documents([Document(page_content="foo")])
docs_with_score = db.similarity_search_with_score("foo")
print(docs_with_score[0])

The output is:

(Document(page_content='foo'), 0.0)

2 - LlamaIndex

LlamaIndex is a simple, flexible data framework for connecting custom data sources to large language models (LLMs).

pgvecto.rs provides a LlamaIndex integration that allows you to retrieve the most similar vectors.

Install dependencies

Some dependencies are required to use the LlamaIndex integration:

pip install llama-index "pgvecto_rs[sdk]"

You could start the postgres instance with pgvecto.rs extension in a docker container:

docker run \
  --name pgvecto-rs-demo \
  -e POSTGRES_PASSWORD=mysecretpassword \
  -p 5432:5432 \
  -d tensorchord/pgvecto-rs:pg16-v0.2.0

Then you can connect to the database using the psql command line tool. The default username is postgres, and the default password is mysecretpassword.

psql postgresql://postgres:mysecretpassword@localhost:5432/postgres

Run the following SQL to ensure the extension is enabled.

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

Create the database and load documents

First, create the text loader and the text splitter to split the text into chunks. We use the markdown file getting-started/overview.md as an example.

import logging
import os
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from pgvecto_rs.sdk import PGVectoRs

URL = "postgresql+psycopg://{username}:{password}@{host}:{port}/{db_name}".format(
    port=os.getenv("DB_PORT", "5432"),
    host=os.getenv("DB_HOST", "localhost"),
    username=os.getenv("DB_USER", "postgres"),
    password=os.getenv("DB_PASS", "mysecretpassword"),
    db_name=os.getenv("DB_NAME", "postgres"),
)

client = PGVectoRs(
    db_url=URL,
    collection_name="example",
    dimension=1536,  # Using OpenAI’s text-embedding-ada-002
)

Then setup the OpenAI API key and load the documents into the database.

import os

os.environ["OPENAI_API_KEY"] = "sk-..."

from llama_index import SimpleDirectoryReader, VectorStoreIndex
from llama_index.vector_stores import PGVectoRsStore

# load documents
documents = SimpleDirectoryReader("./src/getting-started/overview.md").load_data()

# initialize without metadata filter
from llama_index.storage.storage_context import StorageContext

vector_store = PGVectoRsStore(client=client)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents, storage_context=storage_context
)

Query index

Finally, we can retrieve the most similar chunks.

# set Logging to DEBUG for more detailed outputs
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")

3 - Citus

citus is a PostgreSQL extension that turns Postgres into a distributed database. pgvecto.rs works well with citus Here are the steps to enable vector search in

citus is a PostgreSQL extension that turns Postgres into a distributed database. pgvecto.rs works well with citus Here are the steps to enable vector search in a distributed PostgreSQL instance.

citus can works on Single-Node and Multi-Node, both of which are compatible with pgvecto.rs natively.

Single-Node Citus

First, prepare the dockerfile with pgvecto.rs and citus installed and attach them to PostgreSQL.

Details

citus must be placed first in shared_preload_libraries due to its restrictions.

FROM tensorchord/pgvecto-rs-binary:pg16-v0.3.0-amd64 as binary

FROM postgres:16
COPY --from=binary /pgvecto-rs-binary-release.deb /tmp/vectors.deb
RUN apt-get update && apt-get install -y curl
RUN apt-get install -y /tmp/vectors.deb && rm -f /tmp/vectors.deb
RUN curl https://install.citusdata.com/community/deb.sh -o deb.sh
RUN bash deb.sh 
RUN apt-get install -y postgresql-16-citus-12.1

CMD ["postgres", "-c" ,"shared_preload_libraries=citus,vectors.so", "-c", "search_path=\"$user\", public, vectors"]

Then we will build a Docker image called citus-vector:16 locally.

Warning

Don’t use POSTGRES_HOST_AUTH_METHOD: "trust" in production, see citus notes about Increasing Worker Security for detailed authentication configure.

docker build -t citus-vector:16 .
docker run --name citus-single -e POSTGRES_HOST_AUTH_METHOD=trust -p 5432:5432 -d citus-vector:16

To enable the essential extensions, connect our container with psql and execute SQL command.

-- psql -h 127.0.0.1 -d postgres -U postgres -w
CREATE EXTENSION IF NOT EXISTS citus;
CREATE EXTENSION IF NOT EXISTS vectors;

Once the extensions are created, you can create a simple table, then transform it into a distributed table.

CREATE TABLE items (id bigserial, embedding vector(3), category_id bigint, PRIMARY KEY (id, category_id));
SET citus.shard_count = 4;
SELECT create_distributed_table('items', 'category_id');

The index build and vector search is as same as before. Tables and indexes are distributed to all shards/workers automatically.

INSERT INTO items (embedding, category_id)
SELECT '[1, 1, 1]'::vector, FLOOR(random() * 100) FROM generate_series(1, 1000)
UNION ALL
SELECT '[1, 2, 3]'::vector, FLOOR(random() * 100) FROM generate_series(1, 1000)
UNION ALL
SELECT '[3, 2, 1]'::vector, FLOOR(random() * 100) FROM generate_series(1, 1000);
CREATE INDEX ON items USING vectors (embedding vector_cos_ops) WITH (options = "[indexing.hnsw]");
SELECT id FROM items ORDER BY embedding <=> '[1, 1, 1]' LIMIT 100;

Multi-Node Citus

You can use the same Dockerfile and image for Multi-Node citus as for Single-Node Citus. In this example, we will start a cluster of 1 coordinator node and 2 worker nodes with docker compose.

Warning

Don’t use POSTGRES_HOST_AUTH_METHOD: "trust" in production, see citus notes about Increasing Worker Security for detailed authentication configure.

services:
  main:
    container_name: main
    image: 'citus-vector:16'
    ports:
      - "5432:5432"
    environment:
      POSTGRES_HOST_AUTH_METHOD: "trust"
  worker1:
    container_name: worker1
    image: 'citus-vector:16'
    ports:
      - "5431:5432"
    environment:
      POSTGRES_HOST_AUTH_METHOD: "trust"
  worker2:
    container_name: worker2
    image: 'citus-vector:16'
    ports:
      - "5430:5432"
    environment:
      POSTGRES_HOST_AUTH_METHOD: "trust"

Now start all containers with docker-compose. Each container is a citus node, with 1 coordinator(main) and 2 workers.

docker compose up -d

As before, you’ll want to enable the extensions, but make sure you run the command on all nodes.

psql -h 127.0.0.1 -d postgres -U postgres -w -c 'CREATE EXTENSION IF NOT EXISTS citus;CREATE EXTENSION IF NOT EXISTS vectors;'
psql -h 127.0.0.1 -d postgres -U postgres -w -p 5431 -c 'CREATE EXTENSION IF NOT EXISTS citus;CREATE EXTENSION IF NOT EXISTS vectors;'
psql -h 127.0.0.1 -d postgres -U postgres -w -p 5430 -c 'CREATE EXTENSION IF NOT EXISTS citus;CREATE EXTENSION IF NOT EXISTS vectors;'

To inform the coordinator about its workers, connect to the coordinator node with psql and register the two workers.

-- psql -h 127.0.0.1 -d postgres -U postgres -w
SELECT citus_set_coordinator_host('main', 5432);
SELECT * from citus_add_node('worker1', 5432);
SELECT * from citus_add_node('worker2', 5432);

Finally, you can perform vector queries in the same way as in Single-Node Citus.

CREATE TABLE items (id bigserial, embedding vector(3), category_id bigint, PRIMARY KEY (id, category_id));
SET citus.shard_count = 4;
SELECT create_distributed_table('items', 'category_id');
INSERT INTO items (embedding, category_id) 
SELECT '[1, 1, 1]'::vector, FLOOR(random() * 100) FROM generate_series(1, 1000)
UNION ALL
SELECT '[1, 2, 3]'::vector, FLOOR(random() * 100) FROM generate_series(1, 1000)
UNION ALL
SELECT '[3, 2, 1]'::vector, FLOOR(random() * 100) FROM generate_series(1, 1000);
CREATE INDEX ON items USING vectors (embedding vector_cos_ops) WITH (options = "[indexing.hnsw]");
SELECT id FROM items ORDER BY embedding <-> '[1, 1, 1]' LIMIT 100;

Monitor

If you’re using a citus distributed table, the pg_vector_index_stat view on the coordinator will be empty. Since all the indexes are created at the workers actually, we can only inspect the vector status by distributing this view to the workers.

SELECT * from run_command_on_workers($$ SELECT (to_json(array_agg(pg_vector_index_stat .*))) FROM pg_vector_index_stat $$);