Skip to main content

Rockset

Rockset是一个为云端构建的实时搜索和分析数据库。Rockset使用具有高效存储向量嵌入的Converged Index™，以在规模上提供低延迟、高并发的搜索查询。Rockset完全支持元数据过滤，并处理不断更新的流式数据的实时摄取。

本教程演示了如何在LangChain中将Rockset用作向量存储。在开始之前，请确保您可以访问Rockset帐户并且有一个可用的API密钥。立即开始免费试用。

设置您的环境​

  1. 利用Rockset控制台创建一个使用写API作为源的collection。在本教程中，我们创建了一个名为langchain_demo的collection。

    配置以下摄取转换来标记您的嵌入字段并利用性能和存储优化：

    （我们在这个示例中使用了OpenAI的text-embedding-ada-002，其中#length_of_vector_embedding = 1536）

SELECT _input.* EXCEPT(_meta), 
VECTOR_ENFORCE(_input.description_embedding, #length_of_vector_embedding, 'float') as description_embedding 
FROM _input

  1. 创建集合后，使用控制台检索一个API密钥。对于本教程，我们假设您正在使用Oregon(us-west-2)地区。

  2. 安装rockset-python-client以使LangChain能够直接与Rockset通信。

pip install rockset

LangChain教程​

在您自己的Python笔记本中跟随以下步骤生成并存储Rockset中的向量嵌入。 开始使用Rockset搜索与您的搜索查询相似的文档。

1. 定义关键变量​

import os
import rockset

ROCKSET_API_KEY = os.environ.get("ROCKSET_API_KEY") # 验证ROCKSET_API_KEY环境变量
ROCKSET_API_SERVER = rockset.Regions.usw2a1 # 验证Rockset地区
rockset_client = rockset.RocksetClient(ROCKSET_API_SERVER, ROCKSET_API_KEY)

COLLECTION_NAME='langchain_demo'
TEXT_KEY='description'
EMBEDDING_KEY='description_embedding'

2. 准备文档​

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.document_loaders import TextLoader
from langchain.vectorstores import Rockset

loader = TextLoader('../../../state_of_the_union.txt')
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

3. 插入文档​

embeddings = OpenAIEmbeddings() # 验证OPENAI_API_KEY环境变量

docsearch = Rockset(
    client=rockset_client,
    embeddings=embeddings,
    collection_name=COLLECTION_NAME,
    text_key=TEXT_KEY,
    embedding_key=EMBEDDING_KEY,
)

ids=docsearch.add_texts(
    texts=[d.page_content for d in docs],
    metadatas=[d.metadata for d in docs],
)

4. 搜索相似文档​

query = "What did the president say about Ketanji Brown Jackson"
output = docsearch.similarity_search_with_relevance_scores(
    query, 4, Rockset.DistanceFunction.COSINE_SIM
)
print("output length:", len(output))
for d, dist in output:
    print(dist, d.metadata, d.page_content[:20] + '...')

##
# output length: 4
# 0.764990692109871 {'source': '../../../state_of_the_union.txt'} Madam Speaker, Madam...
# 0.7485416901622112 {'source': '../../../state_of_the_union.txt'} And I’m taking robus...
# 0.7468678973398306 {'source': '../../../state_of_the_union.txt'} And so many families...
# 0.7436231261419488 {'source': '../../../state_of_the_union.txt'} Groups of citizens b...

5. 使用过滤器搜索相似文档​

output = docsearch.similarity_search_with_relevance_scores(
    query,
    4,
    Rockset.DistanceFunction.COSINE_SIM,
    where_str="{} NOT LIKE '%citizens%'".format(TEXT_KEY),
)
print("output length:", len(output))
for d, dist in output:
    print(dist, d.metadata, d.page_content[:20] + '...')

##
# output length: 4
# 0.7651359650263554 {'source': '../../../state_of_the_union.txt'} Madam Speaker, Madam...
# 0.7486265516824893 {'source': '../../../state_of_the_union.txt'} And I’m taking robus...
# 0.7469625542348115 {'source': '../../../state_of_the_union.txt'} And so many families...
# 0.7344177777547739 {'source': '../../../state_of_the_union.txt'} We see the unity amo...

6. [可选] 删除插入的文档​

您必须具有与每个文档关联的唯一ID，才能从集合中删除它们。在插入文档时使用Rockset.add_texts()定义ID。否则，Rockset将为每个文档生成一个唯一ID。无论如何，Rockset.add_texts()都会返回插入文档的ID。

要删除这些文档，只需使用Rockset.delete_texts()函数。

docsearch.delete_texts(ids)

总结​

在本教程中，我们成功创建了一个Rockset集合，使用OpenAI嵌入插入了文档，并使用和不使用元数据过滤器搜索了相似的文档。

请关注https://rockset.com/以获取此领域的未来更新。