Skip to main content
Rockset 是一个专为云构建的实时搜索与分析数据库。Rockset 使用 Converged Index™，并配备高效的向量嵌入存储，可大规模提供低延迟、高并发的搜索查询。Rockset 完全支持元数据过滤，并能处理持续更新的流式数据实时摄取。
本 notebook 演示了如何在 LangChain 中将 Rockset 用作向量存储。在开始之前，请确保您可以访问 Rockset 账户并有可用的 API 密钥。立即开始免费试用。 使用此集成需要通过 pip install -qU langchain-community 安装 langchain-community。

​
环境设置

  1. 利用 Rockset 控制台创建一个以 Write API 作为数据源的集合。在本演示中，我们创建一个名为 langchain_demo 的集合。 配置以下摄取转换以标记您的嵌入字段，并充分利用性能和存储优化： （本示例使用 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 密钥。本 notebook 假设您使用 Oregon(us-west-2) 区域。
  2. 安装 rockset-python-client，使 LangChain 能够直接与 Rockset 通信。 
pip install -qU  rockset

​
LangChain 教程

在您自己的 Python notebook 中跟随操作，在 Rockset 中生成并存储向量嵌入。 然后使用 Rockset 搜索与您查询相似的文档。

​
1. 定义关键变量

import os

import rockset

ROCKSET_API_KEY = os.environ.get(
    "ROCKSET_API_KEY"
)  # Verify ROCKSET_API_KEY environment variable
ROCKSET_API_SERVER = rockset.Regions.usw2a1  # Verify Rockset region
rockset_client = rockset.RocksetClient(ROCKSET_API_SERVER, ROCKSET_API_KEY)

COLLECTION_NAME = "langchain_demo"
TEXT_KEY = "description"
EMBEDDING_KEY = "description_embedding"

​
2. 准备文档

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import Rockset
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

loader = TextLoader("../../how_to/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()  # Verify OPENAI_API_KEY environment variable

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 嵌入插入了文档，并在带/不带元数据过滤器的情况下搜索了相似文档。 关注 rockset.com/ 获取该领域的未来更新。
在 GitHub 上编辑此页面 或 提交 issue。
将这些文档连接 到 Claude、VSCode 等，通过 MCP 获取实时答案。