Skip to main content
Open In ColabOpen on GitHub

Databricks 向量搜索

Databricks Vector Search 是一个无服务器相似搜索引擎，允许你将数据的向量表示及其元数据存储在向量数据库中。通过 Vector Search，你可以从由 Unity Catalog 管理的 Delta 表中创建自动更新的向量搜索索引，并使用简单的 API 查询这些索引来返回最相似的向量。

这个笔记本展示了如何在 Databricks 向量搜索中使用 LangChain。

设置​

要访问Databricks模型，您需要创建一个Databricks账户、设置凭据（如果您在Databricks工作区之外，请仅在此步骤操作）并安装所需的包。

Credentials（仅如果您在Databricks之外）​

如果您在Databricks中运行LangChain应用，请跳过此步骤。

否则，您需要手动设置 Databricks 工作空间主机名和个人访问令牌环境变量分别为 DATABRICKS_HOST 和 DATABRICKS_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: "
    )

安装​

The LangChain Databricks 整合功能位于 databricks-langchain 包中。

%pip install -qU databricks-langchain

创建向量搜索端点和索引（如果还没有创建的话）​

在本部分，我们将使用客户端SDK创建一个Databricks向量搜索端点和索引。

如果您已经有一个端点和一个索引，可以直接跳过该部分并直接进入“初始化”部分。

首先，实例化 Databricks VectorSearch 客户端：

from databricks.vector_search.client import VectorSearchClient

client = VectorSearchClient()

接下来，我们将创建一个新的向量搜索端点。

endpoint_name = "<your-endpoint-name>"

client.create_endpoint(name=endpoint_name, endpoint_type="STANDARD")

最后，我们将创建一个可以在端点上查询的索引。Databricks Vector Search中有两种类型的索引，而DatabricksVectorSearch类支持这两种用例。

  • Delta 同步索引会自动与源 Delta 表进行同步，并在 Delta 表下的数据发生变化时，自动和增量地更新索引。

  • 直接向量访问索引支持直接读取和写入向量及元数据。用户需使用REST API或Python SDK自行更新此表。

对于增量同步索引，您可以选择使用Databricks托管的嵌入式表示或自托管的嵌入式表示（通过LangChain嵌入类实现）。</span>

The following code creates a 直接访问索引。请参阅 Databricks 文档 以创建其他类型的索引。

index_name = "<your-index-name>"  # Format: "<catalog>.<schema>.<index-name>"

index = client.create_direct_access_index(
    endpoint_name=endpoint_name,
    index_name=index_name,
    primary_key="id",
    # Dimension of the embeddings. Please change according to the embedding model you are using.
    embedding_dimension=3072,
    # A column to store the embedding vectors for the text data
    embedding_vector_column="text_vector",
    schema={
        "id": "string",
        "text": "string",
        "text_vector": "array<float>",
        # Optional metadata columns
        "source": "string",
    },
)

index.describe()

Instantiation​

使用Databricks托管的嵌入式向量还是自行管理的嵌入式向量（即您选择的LangChain Embeddings对象）,DatabricksVectorSearch 的实例化方式会有所不同。

如果您正在使用与Databricks管理的嵌入相关的delta-sync索引：

from databricks_langchain import DatabricksVectorSearch

vector_store = DatabricksVectorSearch(
    endpoint=endpoint_name,
    index_name=index_name,
)

如果您使用直接访问索引或自管理嵌入的增量同步索引，您还需要在源表中提供用于嵌入的嵌入模型和文本列：

pip install -qU langchain-openai
import getpass
import os

if not os.environ.get("OPENAI_API_KEY"):
  os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

from langchain_openai import OpenAIEmbeddings

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
vector_store = DatabricksVectorSearch(
    endpoint=endpoint_name,
    index_name=index_name,
    embedding=embeddings,
    # The column name in the index that contains the text data to be embedded
    text_column="document_content",
)

管理向量存储​

添加项到向量存储​

注意：仅对直接访问索引支持通过add_documents方法向向量存储中添加项目。

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"])
API 参考:文档
['1', '2', '3']

删除向量存储中的项​

注：通过delete方法从向量存储中删除项仅支持直接访问索引。

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参考：https://api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.DatabricksVectorSearch