这是本节的多页打印视图。 .

生态集成

将 pgvecto.rs 与 LangChain、LlamaIndex 以及 Citus 分布式 PostgreSQL 集成。

1: LangChain
2: LlamaIndex
3: Citus

1 - LangChain

本文介绍如何将 pgvecto.rs 与 LangChain 集成。

本文介绍如何将 pgvecto.rs 与 LangChain 集成。

LangChain 是一个用于开发由语言模型驱动的应用程序的框架。它提供了一套工具，支持在生产环境中微调、部署和使用语言模型。

pgvecto.rs 提供了 LangChain 集成，支持在 LangChain 中检索最相似的向量。

安装依赖

使用 LangChain 集成需要安装以下依赖：

pip install langchain langchain-openai pgvecto-rs

你可以在 Docker 容器中启动带有 pgvecto.rs 扩展的 Postgres 实例：

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

然后可以使用 psql 命令行工具连接到数据库。默认用户名为 postgres，默认密码为 mysecretpassword。

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

执行以下 SQL 以确保扩展已启用：

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

创建数据库并加载文档

下面展示如何在 LangChain 中使用 pgvecto.rs 检索最相似的向量。

首先，你需要创建文本加载器（TextLoader）和文本切分器（CharacterTextSplitter），将文本切分为多个片段（chunk）。这里以 Markdown 文件 getting-started/overview.md 为例。

## 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()

接下来，创建 PGVecto_rs 实例并将文档加载到数据库中。

## 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,
)

查询索引

最后，我们可以在 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)

输出结果如下：

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.
...

从现有数据库初始化

在上面的示例中，我们从零开始创建了一个向量存储。但在很多时候，我们希望连接并使用已有的向量存储，此时可以直接对其进行初始化：

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

随后，我们可以向存储中添加向量并进行查询：

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

输出结果如下：

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

2 - LlamaIndex

LlamaIndex 是一个简单、灵活的数据框架，用于将自定义数据源连接到大语言模型（LLM）。

pgvecto.rs 提供了 LlamaIndex 集成，支持检索最相似的向量。

安装依赖

使用 LlamaIndex 集成需要安装以下依赖：

pip install llama-index "pgvecto_rs[sdk]"

你可以在 Docker 容器中启动带有 pgvecto.rs 扩展的 Postgres 实例：

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

然后可以使用 psql 命令行工具连接到数据库。默认用户名为 postgres，默认密码为 mysecretpassword。

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

执行以下 SQL 以确保扩展已启用：

DROP EXTENSION IF EXISTS vectors;
CREATE EXTENSION vectors;

创建数据库并加载文档

首先，创建数据库连接客户端。这里同样以 Markdown 文件 getting-started/overview.md 为例。

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
)

接着配置 OpenAI API 密钥，并将文档加载到数据库中。

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
)

查询索引

最后，我们可以检索最相似的文档片段。

# 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 是将 Postgres 转化为分布式数据库的 PostgreSQL 扩展。pgvecto.rs 可与 Citus 良好配合，本文介绍在分布式 PostgreSQL 实例中启用向量检索的步骤。

citus 是一个将 Postgres 转化为分布式数据库的 PostgreSQL 扩展。pgvecto.rs 能够与 citus 良好协作。下面介绍如何在分布式 PostgreSQL 实例中启用向量检索。

citus 支持在单节点（Single-Node）和多节点（Multi-Node）模式下运行，这两种模式均原生兼容 pgvecto.rs。

单节点 Citus

首先，准备一个同时安装了 pgvecto.rs 和 citus 的 Dockerfile，并将它们加载到 PostgreSQL 中。

详情

由于 citus 的自身限制，在 shared_preload_libraries 中必须将 citus 放在第一位。

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"]

接着在本地构建名为 citus-vector:16 的 Docker 镜像并启动容器。

警告

请勿在生产环境中使用 POSTGRES_HOST_AUTH_METHOD: "trust"。详细的认证配置请参阅 citus 文档中关于提高 Worker 安全性（Increasing Worker Security）的说明。

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

要启用所需扩展，请使用 psql 连接到容器并执行以下 SQL 命令：

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

扩展创建完成后，你可以先创建一张普通表，然后将其转换为分布式表（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');

之后的索引构建和向量检索操作与单机版完全一致。表和索引会自动分发到所有分片（Shard）/ 工作节点（Worker）上。

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;

多节点 Citus

你可以复用与单节点 Citus 相同的 Dockerfile 和镜像来部署多节点 citus。在本示例中，我们将使用 Docker Compose 启动一个包含 1 个协调器节点（Coordinator）和 2 个工作节点（Worker）的集群。

警告

请勿在生产环境中使用 POSTGRES_HOST_AUTH_METHOD: "trust"。详细的认证配置请参阅 citus 文档中关于提高 Worker 安全性（Increasing Worker Security）的说明。

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"

使用 Docker Compose 启动所有容器。每个容器对应一个 citus 节点，包括 1 个协调器节点（main）和 2 个工作节点。

docker compose up -d

与之前一样，需要启用相关扩展，但请确保在所有节点上都执行该命令：

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;'

为了让协调器知晓工作节点的信息，请使用 psql 连接到协调器节点，并注册这两个工作节点：

-- 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);

最后，你可以按照与单节点 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;

监控

如果你使用的是 citus 分布式表，协调器节点上的 pg_vector_index_stat 视图将为空。因为所有索引实际上都创建在工作节点上，所以我们需要将查询分发到各个工作节点来查看向量索引状态：

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