Skip to main content
Databricks Vector Search 是一款无服务器相似度搜索引擎,允许您将数据(包括元数据)的向量表示存储在向量数据库中。借助 Vector Search,您可以从 Unity Catalog 管理的 Delta 表创建自动更新的向量搜索索引,并通过简单的 API 查询返回最相似的向量。 本 notebook 展示如何在 Databricks Vector Search 中使用 LangChain。

设置

要访问 Databricks 模型,您需要创建 Databricks 账号,设置凭证(仅在 Databricks 工作空间外时需要),并安装所需包。

凭证(仅在 Databricks 工作空间外时需要)

如果您在 Databricks 内运行 LangChain 应用,可以跳过此步骤。 否则,您需要手动将 Databricks 工作空间主机名和个人访问令牌分别设置到 DATABRICKS_HOSTDATABRICKS_TOKEN 环境变量中。有关如何获取访问令牌,请参阅认证文档
import getpass
import os

os.environ["DATABRICKS_HOST"] = "https://your-databricks-workspace"
if "DATABRICKS_TOKEN" not in os.environ:
    os.environ["DATABRICKS_TOKEN"] = getpass.getpass(
        "Enter your Databricks access token: "
    )

安装

LangChain Databricks 集成位于 databricks-langchain 包中。
pip install -qU databricks-langchain

创建向量搜索端点和索引(如果尚未创建)

本节将使用客户端 SDK 创建 Databricks Vector Search 端点和索引。 如果您已有端点和索引,可以跳过本节直接进入”实例化”部分。 首先,实例化 Databricks VectorSearch 客户端:
from databricks.vector_search.client import VectorSearchClient

client = VectorSearchClient()
接下来,创建一个新的 VectorSearch 端点。
endpoint_name = "<your-endpoint-name>"

client.create_endpoint(name=endpoint_name, endpoint_type="STANDARD")
最后,在该端点上创建可查询的索引。Databricks Vector Search 有两种索引类型,DatabricksVectorSearch 类均支持:
  • Delta Sync 索引 自动与源 Delta 表同步,随着 Delta 表中数据的变化自动增量更新索引。
  • Direct Vector Access 索引 支持对向量和元数据的直接读写。用户负责通过 REST API 或 Python SDK 更新该表。
对于 delta-sync 索引,您可以选择使用 Databricks 托管嵌入或自托管嵌入(通过 LangChain 嵌入类)。 以下代码创建 direct-access 索引。关于其他类型索引的创建说明,请参阅 Databricks 文档
index_name = "<your-index-name>"  # 格式:"<catalog>.<schema>.<index-name>"

index = client.create_direct_access_index(
    endpoint_name=endpoint_name,
    index_name=index_name,
    primary_key="id",
    # 嵌入维度,请根据您使用的嵌入模型修改
    embedding_dimension=3072,
    # 用于存储文本数据嵌入向量的列
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "text": "string",
        "text_vector": "array<float>",
        # 可选的元数据列
        "source": "string",
    },
)

index.describe()

实例化

DatabricksVectorSearch 的实例化方式因索引使用 Databricks 托管嵌入还是自托管嵌入(即您选择的 LangChain Embeddings 对象)而略有不同。 如果您使用的是带 Databricks 托管嵌入的 delta-sync 索引:
from databricks_langchain import DatabricksVectorSearch

vector_store = DatabricksVectorSearch(
    endpoint=endpoint_name,
    index_name=index_name,
)
如果您使用的是 direct-access 索引或带自托管嵌入的 delta-sync 索引,还需要提供嵌入模型以及源表中用于嵌入的文本列:
# | output: false
# | echo: false
from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
vector_store = DatabricksVectorSearch(
    endpoint=endpoint_name,
    index_name=index_name,
    embedding=embeddings,
    # 索引中包含待嵌入文本数据的列名
    text_column="document_content",
)

管理向量存储

向向量存储添加条目

注意:通过 add_documents 方法向向量存储添加条目仅支持 direct-access 索引。
from langchain_core.documents import Document

document_1 = Document(page_content="foo", metadata={"source": "https://example.com"})

document_2 = Document(page_content="bar", metadata={"source": "https://example.com"})

document_3 = Document(page_content="baz", metadata={"source": "https://example.com"})

documents = [document_1, document_2, document_3]

vector_store.add_documents(documents=documents, ids=["1", "2", "3"])
['1', '2', '3']

从向量存储删除条目

注意:通过 delete 方法从向量存储删除条目仅支持 direct-access 索引。
vector_store.delete(ids=["3"])
True

查询向量存储

创建向量存储并添加相关文档后,通常需要在运行链或智能体时对其进行查询。

直接查询

执行简单相似度搜索的方法如下:
results = vector_store.similarity_search(
    query="thud", k=1, filter={"source": "https://example.com"}
)
for doc in results:
    print(f"* {doc.page_content} [{doc.metadata}]")
* foo [{'id': '1'}]
注意:默认情况下,相似度搜索仅返回主键和文本列。如果想要检索文档关联的自定义元数据,请在初始化向量存储时通过 columns 参数传入额外的列名。
vector_store = DatabricksVectorSearch(
    endpoint=endpoint_name,
    index_name=index_name,
    embedding=embeddings,
    text_column="text",
    columns=["source"],
)

results = vector_store.similarity_search(query="thud", k=1)
for doc in results:
    print(f"* {doc.page_content} [{doc.metadata}]")
* foo [{'source': 'https://example.com', 'id': '1'}]
如果要执行相似度搜索并获取对应分数,可以运行:
results = vector_store.similarity_search_with_score(
    query="thud", k=1, filter={"source": "https://example.com"}
)
for doc, score in results:
    print(f"* [SIM={score:3f}] {doc.page_content} [{doc.metadata}]")
* [SIM=0.414035] foo [{'source': 'https://example.com', 'id': '1'}]

转换为检索器进行查询

您还可以将向量存储转换为检索器,以便在链中更方便地使用。
retriever = vector_store.as_retriever(search_type="mmr", search_kwargs={"k": 1})
retriever.invoke("thud")
[Document(metadata={'source': 'https://example.com', 'id': '1'}, page_content='foo')]

用于检索增强生成

有关如何将此向量存储用于检索增强生成(RAG)的指南,请参阅以下章节:

API 参考

有关 DatabricksVectorSearch 所有功能和配置的详细文档,请前往 API 参考:api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.DatabricksVectorSearch