使用pgvecto.rs实现向量存储与查询

最新推荐文章于 2024-07-12 19:30:19 发布

ppoojjj

最新推荐文章于 2024-07-12 19:30:19 发布

阅读量287

点赞数 4

文章标签： python

本文链接：https://blog.csdn.net/ppoojjj/article/details/140366805

版权

在本文中，我们将介绍如何使用pgvecto.rs来实现向量存储与查询。我们将详细解释如何安装依赖、配置pgvecto.rs客户端、加载文档并建立向量存储索引，以及如何进行查询。

安装依赖

首先，我们需要安装所需的Python依赖包：

%pip install llama-index-vector-stores-pgvecto-rs
%pip install llama-index "pgvecto_rs[sdk]"

接下来，启动pgvecto.rs服务器：

!docker run --name pgvecto-rs-demo -e POSTGRES_PASSWORD=mysecretpassword -p 5432:5432 -d tensorchord/pgvecto-rs:latest

配置日志记录

为了方便调试和查看详细日志，我们需要配置日志记录器：

import logging
import os
import sys

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

创建pgvecto_rs客户端

我们可以通过以下代码创建一个pgvecto.rs客户端：

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,  # 使用OpenAI的text-embedding-ada-002
)

配置OpenAI

接下来，我们需要设置OpenAI的API密钥：

import os

os.environ["OPENAI_API_KEY"] = "sk-..."  # 请替换为你自己的API密钥

加载文档，构建PGVectoRsStore和VectorStoreIndex

我们将下载一个示例文档，并加载到我们的向量存储中：

from IPython.display import Markdown, display
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.vector_stores.pgvecto_rs import PGVectoRsStore

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

# 加载文档
documents = SimpleDirectoryReader("./data/paul_graham").load_data()

# 初始化向量存储（不使用元数据过滤）
from llama_index.core 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
)

查询索引

最后，我们可以通过以下代码进行查询：

# 设置日志级别为DEBUG，以便查看详细输出
query_engine = index.as_query_engine()
response = query_engine.query("作者在成长过程中做了什么？")

import IPython.display as display
display(Markdown(f"<b>{response}</b>"))