DatabricksVectorSearch
Databricks 向量搜索 是一个无服务器的相似性搜索引擎，允许您将数据的向量表示（包括元数据）存储在向量数据库中。通过向量搜索，您可以从 Unity Catalog 管理的 Delta 表中创建自动更新的向量搜索索引，并使用简单的 API 查询它们以返回最相似的向量。
本笔记本展示了如何在 Databricks Vector Search 中使用 LangChain。
设置
要访问 Databricks 模型，你需要创建一个 Databricks 账户、设置凭据（仅当你在 Databricks 工作区之外时），并安装所需的包。
凭据（仅当您在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: "
)
安装
LangChain Databricks 集成位于 databricks-langchain 包中。
%pip install -qU databricks-langchain
创建向量搜索端点和索引（如果您尚未创建）
在本节中，我们将使用客户端 SDK 创建一个 Databricks 向量搜索端点和索引。
如果你已经有端点和索引，可以跳过该部分，直接进入“实例化”部分。
首先，实例化 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 同步索引 会自动与源 Delta 表同步，当 Delta 表中的基础数据发生变化时，会自动且增量地更新索引。
-
直接向量访问索引 支持对向量和元数据的直接读写。用户需通过 REST API 或 Python SDK 更新此表。
对于 delta-sync 索引，您可以选择使用 Databricks 管理的嵌入（embeddings）或自行管理的嵌入（通过 LangChain 嵌入类）。
以下代码创建一个直接访问索引。请参考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()
实例化
DatabricksVectorSearch 的实例化方式会有所不同，具体取决于您的索引是使用 Databricks 管理的嵌入（embeddings）还是自管理的嵌入（即您选择的 LangChain 嵌入对象）。
如果你使用的是与 Databricks 管理的嵌入向量一起使用的增量同步索引：
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"])
['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