Databricks vector search 集成

使用 LangChain Python 集成 Databricks vector search 向量存储。

Databricks Vector Search is a serverless similarity search engine that allows you to store a vector representation of your data, including metadata, in a vector database. With Vector Search, you can create auto-updating vector search indexes from Delta tables managed by Unity Catalog and query them with a simple API to return the most similar vectors.

This notebook shows how to use LangChain with Databricks Vector Search.

设置

To access Databricks models you’ll need to create a Databricks account, set up credentials (only if you are outside Databricks workspace), and install required packages.

Credentials (only if you are outside databricks)

If you are running LangChain app inside Databricks, you can skip this step.

Otherwise, you need manually set the Databricks workspace hostname and personal access token to DATABRICKS_HOST and DATABRICKS_TOKEN environment variables, respectively. See Authentication Documentation for how to get an access 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 integration lives in the databricks-langchain package.

pip install -qU databricks-langchain

Create a vector search endpoint and index (if you haven’t already)

In this section, we will create a Databricks Vector Search endpoint and an index using the client SDK.

If you already have an endpoint and an index, you can skip the section and go straight to “Instantiation” section.

First, instantiate the Databricks VectorSearch client:

from databricks.vector_search.client import VectorSearchClient

client = VectorSearchClient()

Next, we will create a new VectorSearch endpoint.

endpoint_name = "<your-endpoint-name>"

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

Lastly, we will create an index that can be queried on the endpoint. There are two types of indexes in Databricks Vector Search and the DatabricksVectorSearch class support both use cases.

• Delta Sync Index automatically syncs with a source Delta Table, automatically and incrementally updating the index as the underlying data in the Delta Table changes.

• Direct Vector Access Index supports direct read and write of vectors and metadata. The user is responsible for updating this table using the REST API or the Python SDK.

Also for delta-sync index, you can choose to use Databricks-managed embeddings or self-managed embeddings (via LangChain embeddings classes).

The following code creates a direct-access index. Please refer to the Databricks documentation for the instruction to create the other type of indexes.

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

实例化

The instantiation of DatabricksVectorSearch is a bit different depending on whether your index uses Databricks-managed embeddings or self-managed embeddings i.e. LangChain Embeddings object of your choice.

If you are using a delta-sync index with Databricks-managed embeddings:

from databricks_langchain import DatabricksVectorSearch

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

If you are using a direct-access index or a delta-sync index with self-managed embeddings, you also need to provide the embedding model and text column in your source table to use for the embeddings:

# | 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,
    # The column name in the index that contains the text data to be embedded
    text_column="document_content",
)

管理向量存储

向向量存储添加项目

Note: Adding items to vector store via add_documents method is only supported for a direct-access index.

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']

从向量存储删除项目

Note: Deleting items to vector store via delete method is only supported for a direct-access index.

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'}]

Note: By default, similarity search only returns the primary key and text column. If you want to retrieve the custom metadata associated with the document, pass the additional columns in the columns parameter when initializing the vector store.

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) 的指南，请参阅以下部分：

• Tutorials
• How-to: Question and answer with RAG
• Retrieval conceptual docs

---

API 参考

For detailed documentation of all DatabricksVectorSearch features and configurations head to the API reference: api-docs.databricks.com/python/databricks-ai-bridge/latest/databricks_langchain.html#databricks_langchain.DatabricksVectorSearch

---

连接这些文档到 Claude、VSCode 等工具，通过 MCP 获取实时答案。

在 GitHub 上编辑此页面或提交 issue。