如何使用Rockset进行实时向量搜索:LangChain集成指南

71 阅读2分钟
# 引言

在现代数据处理领域,实时搜索和分析数据库是一个非常关键的组件。Rockset是一款为云端设计的实时搜索和分析数据库,能够有效处理低延迟、高并发的搜索查询。本篇文章将介绍如何利用Rockset作为LangChain中的向量存储,帮助开发者高效地进行文档相似性搜索。

# 主要内容

## 1. 环境设置

首先,你需要在Rockset控制台中创建一个集合(collection),这里我们命名为`langchain_demo`。请确保你已经拥有一个Rockset账户和API密钥。可以通过以下步骤进行设置:

- 使用Write API创建集合。
- 配置数据摄取转换以优化性能和存储。

设置了集合后,你可以在控制台中获取API密钥。

### 2. 安装依赖

在开始之前,请确保安装必要的Python包:

```bash
pip install -qU langchain-community rockset

3. 定义关键变量

在你的Python环境中定义以下变量:

import os
import rockset

ROCKSET_API_KEY = os.environ.get("ROCKSET_API_KEY")  # 请确认API密钥已设置为环境变量
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"

4. 准备文档

使用LangChain内建的工具来加载和处理文档:

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)

5. 插入文档

创建嵌入向量并插入文档:

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

6. 搜索相似文档

可以在Rockset中进行相似性搜索:

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

常见问题和解决方案

  • 网络访问问题:由于某些地区可能存在网络限制,建议使用API代理服务来提高访问稳定性。
  • 性能优化:合理配置数据摄取转换以优化存储和查询性能。

总结和进一步学习资源

通过本文的介绍,我们成功创建了Rockset集合,插入了带有OpenAI嵌入的文档,并进行了相似性搜索。更多相关细节请参阅Rockset文档

参考资料

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---