探索在LangChain中使用Rockset进行实时向量搜索

最新推荐文章于 2024-11-08 17:19:46 发布

nseejrukjhad

最新推荐文章于 2024-11-08 17:19:46 发布

阅读量264

点赞数 5

文章标签： langchain 服务器算法 python

本文链接：https://blog.csdn.net/nseejrukjhad/article/details/142747614

版权

探索在LangChain中使用Rockset进行实时向量搜索

引言

在现代数据驱动的世界中，实时搜索和分析变得越来越重要。Rockset是一个为云构建的实时搜索和分析数据库，它能够快速处理高并发的搜索查询。这篇文章将介绍如何在LangChain中使用Rockset作为向量存储，并提供详细的代码示例和实用的见解。

主要内容

安装和环境设置

首先，你需要安装langchain-community库：

pip install -qU langchain-community

接下来，通过Rockset控制台创建一个名为langchain_demo的集合，并配置以下数据导入转换：

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

还需要安装rockset-python-client以便LangChain与Rockset直接通信：

pip install --upgrade --quiet rockset

设置关键变量

以下是一些关键变量的定义：

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"

代码示例

准备和插入文档

使用以下代码准备并插入文档：

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)

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

搜索类似文档

可以通过以下方式进行文档搜索：

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