探索Databricks Vector Search:构建自我查询检索器的实战指南

93 阅读3分钟
# 探索Databricks Vector Search:构建自我查询检索器的实战指南

## 引言

在现代数据驱动的世界中,快速有效的相似性搜索成为关键任务。Databricks Vector Search提供了一种无服务器的解决方案,能够在向量数据库中存储数据的向量表示及其元数据。本篇文章将介绍如何使用Databricks Vector Search创建并使用自我查询检索器(SelfQueryRetriever),以便从Delta表中创建自动更新的向量搜索索引。

## 主要内容

### 1. 环境准备

首先,我们需要安装所需的Python包。请确保安装 `lark` 和其他相关包:

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

注意:安装后可能需要重启内核以生效。

2. 设置API密钥

我们将使用OpenAI的API,因此需要获取API密钥。为了提高访问稳定性,建议使用API代理服务:

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")  # 使用API代理服务提高访问稳定性
databricks_host = getpass.getpass("Databricks host:")
databricks_token = getpass.getpass("Databricks token:")

3. 创建向量存储索引

接下来,我们需要创建一个Databricks Vector Store Index,并使用电影摘要数据进行初始化:

from databricks.vector_search.client import VectorSearchClient
from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()
emb_dim = len(embeddings.embed_query("hello"))

vsc = VectorSearchClient(
    workspace_url=databricks_host, personal_access_token=databricks_token
)
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="udhay_demo.10x.demo_index",
    primary_key="id",
    embedding_dimension=emb_dim,
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "page_content": "string",
        "year": "int",
        "rating": "float",
        "genre": "string",
        "text_vector": "array<float>",
    },
)

4. 添加文档到向量存储

我们可以通过以下代码将电影摘要文档添加到向量存储中:

from langchain_core.documents import Document

docs = [
    Document(page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose", metadata={"id": 1, "year": 1993, "rating": 7.7, "genre": "action"}),
    # 更多文档
]

from langchain_community.vectorstores import DatabricksVectorSearch

vector_store = DatabricksVectorSearch(
    index,
    text_column="page_content",
    embedding=embeddings,
    columns=["year", "rating", "genre"],
)

vector_store.add_documents(docs)

5. 创建自我查询检索器

我们现在可以设置自我查询检索器,通过描述元数据字段和文档内容来实现智能检索:

from langchain.chains.query_constructor.base import AttributeInfo
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain_openai import OpenAI

metadata_field_info = [
    AttributeInfo(name="genre", description="The genre of the movie", type="string"),
    AttributeInfo(name="year", description="The year the movie was released", type="integer"),
    AttributeInfo(name="rating", description="A 1-10 rating for the movie", type="float"),
]

document_content_description = "Brief summary of a movie"
llm = OpenAI(temperature=0)

retriever = SelfQueryRetriever.from_llm(
    llm, vector_store, document_content_description, metadata_field_info, verbose=True
)

代码示例

使用检索器进行查询:

# 按照特定查询条件进行检索
retriever.invoke("What are some movies about dinosaurs")

我们也可以指定过滤条件:

# 查找评分高于9的电影
retriever.invoke("What are some highly rated movies (above 9)?")

常见问题和解决方案

1. 网络连接问题

因为某些地区网络限制,建议使用API代理服务改善访问稳定性。

2. 数据更新

确保定期更新Delta表以保持搜索结果的准确性。

总结和进一步学习资源

通过Databricks Vector Search和自我查询检索器,开发者可以方便地实现向量相似性搜索,适用于多种应用场景。更多详细信息可以参考以下资源:

参考资料

  1. Databricks官方文档
  2. OpenAI API文档
  3. Langchain核心库指南

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

---END---