将代理连接到非结构化数据

AI 代理通常需要查询非结构化数据（如文档集合、知识库或文本语料库）来回答问题并提供上下文感知响应。

Databricks 提供了多种方法，用于将代理连接到矢量搜索索引和外部矢量存储中的非结构化数据。 使用预配置的 MCP 服务器立即访问 Databricks 矢量搜索索引、使用 AI Bridge 包在本地开发检索器工具，或为专用工作流生成自定义检索器函数。

使用 MCP 查询 Databricks 矢量搜索索引

如果代理需要查询 Databricks 矢量搜索索引，请使用 Databricks 托管的 MCP 服务器。 在开始之前，请使用 Databricks 托管的嵌入创建矢量搜索索引。 请参阅 创建矢量搜索终结点和索引。

矢量搜索的托管 MCP URL 为： https://<workspace-hostname>/api/2.0/mcp/vector-search/{catalog}/{schema}/{index_name}。

以下示例演示如何将代理连接到矢量搜索索引。 替换<catalog>、<schema>和<index-name>为您的矢量搜索索引的名称。

OpenAI 代理 SDK （应用）

from agents import Agent, Runner
from databricks.sdk import WorkspaceClient
from databricks_openai.agents import McpServer

workspace_client = WorkspaceClient()

async with McpServer.from_vector_search(
    catalog="<catalog>",
    schema="<schema>",
    index_name="<index-name>",
    workspace_client=workspace_client,
    name="vector-search",
) as vs_server:
    agent = Agent(
        name="Research assistant",
        instructions="You are a research assistant. Use the vector search tool to find relevant documents and answer questions.",
        model="databricks-claude-sonnet-4-5",
        mcp_servers=[vs_server],
    )
    result = await Runner.run(agent, "What is the return policy?")
    print(result.final_output)

在以下项 databricks.yml中授予应用对矢量搜索索引的访问权限：

resources:
  apps:
    my_agent_app:
      resources:
        - name: 'my_vector_index'
          uc_securable:
            securable_full_name: '<catalog>.<schema>.<index-name>'
            securable_type: 'TABLE'
            permission: 'SELECT'

LangGraph （应用）

from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient
from langgraph.prebuilt import create_react_agent

workspace_client = WorkspaceClient()
host = workspace_client.config.host

mcp_client = DatabricksMultiServerMCPClient([
    DatabricksMCPServer(
        name="vector-search",
        url=f"{host}/api/2.0/mcp/vector-search/<catalog>/<schema>/<index-name>",
        workspace_client=workspace_client,
    ),
])

async with mcp_client:
    tools = await mcp_client.get_tools()
    agent = create_react_agent(
        ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
        tools=tools,
    )
    result = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "What is the return policy?"}]}
    )
    print(result["messages"][-1].content)

在以下项 databricks.yml中授予应用对矢量搜索索引的访问权限：

resources:
  apps:
    my_agent_app:
      resources:
        - name: 'my_vector_index'
          uc_securable:
            securable_full_name: '<catalog>.<schema>.<index-name>'
            securable_type: 'TABLE'
            permission: 'SELECT'

模型服务

from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient
import mlflow

workspace_client = WorkspaceClient()
host = workspace_client.config.host

# Connect to the Vector Search MCP server
mcp_client = DatabricksMCPClient(
    server_url=f"{host}/api/2.0/mcp/vector-search/<catalog>/<schema>/<index-name>",
    workspace_client=workspace_client,
)

# List available tools from the Vector Search index
tools = mcp_client.list_tools()

# Log the agent with the required resources for deployment
mlflow.pyfunc.log_model(
    "agent",
    python_model=my_agent,
    resources=mcp_client.get_databricks_resources(),
)

若要部署代理，请参阅部署用于生成 AI 应用程序的代理（模型服务）。 有关使用 MCP 资源的日志记录代理的详细信息，请参阅 使用 Databricks 托管 MCP 服务器。

其他方法

%pip install --upgrade databricks-langchain

from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks

# Initialize the retriever tool.
vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation."
)

# Run a query against the vector search index locally for testing
vs_tool.invoke("Databricks Agent Framework?")

# Bind the retriever tool to your Langchain LLM of choice
llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([vs_tool])

# Chat with your LLM to test the tool calling functionality
llm_with_tools.invoke("Based on the Databricks documentation, what is Databricks Agent Framework?")

from databricks_langchain import VectorSearchRetrieverTool
from databricks_langchain import DatabricksEmbeddings

embedding_model = DatabricksEmbeddings(
    endpoint="databricks-bge-large-en",
)

vs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
  num_results=5, # Max number of documents to return
  columns=["primary_key", "text_column"], # List of columns to include in the search
  filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
  query_type="ANN", # Query type ("ANN" or "HYBRID").
  tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
  tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
  text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
  embedding=embedding_model # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

%pip install --upgrade databricks-openai

from databricks_openai import VectorSearchRetrieverTool
from openai import OpenAI
import json

# Initialize OpenAI client
client = OpenAI(api_key=<your_API_key>)

# Initialize the retriever tool
dbvs_tool = VectorSearchRetrieverTool(
  index_name="catalog.schema.my_databricks_docs_index",
  tool_name="databricks_docs_retriever",
  tool_description="Retrieves information about Databricks products from official Databricks documentation"
)

messages = [
  {"role": "system", "content": "You are a helpful assistant."},
  {
    "role": "user",
    "content": "Using the Databricks documentation, answer what is Spark?"
  }
]
first_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

# Execute function code and parse the model's response and handle function calls.
tool_call = first_response.choices[0].message.tool_calls[0]
args = json.loads(tool_call.function.arguments)
result = dbvs_tool.execute(query=args["query"])  # For self-managed embeddings, optionally pass in openai_client=client

# Supply model with results – so it can incorporate them into its final response.
messages.append(first_response.choices[0].message)
messages.append({
  "role": "tool",
  "tool_call_id": tool_call.id,
  "content": json.dumps(result)
})
second_response = client.chat.completions.create(
  model="gpt-4o",
  messages=messages,
  tools=[dbvs_tool.tool]
)

from databricks_openai import VectorSearchRetrieverTool

vs_tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.index_name", # Index name in the format 'catalog.schema.index'
    num_results=5, # Max number of documents to return
    columns=["primary_key", "text_column"], # List of columns to include in the search
    filters={"text_column LIKE": "Databricks"}, # Filters to apply to the query
    query_type="ANN", # Query type ("ANN" or "HYBRID").
    tool_name="name of the tool", # Used by the LLM to understand the purpose of the tool
    tool_description="Purpose of the tool", # Used by the LLM to understand the purpose of the tool
    text_column="text_column", # Specify text column for embeddings. Required for direct-access index or delta-sync index with self-managed embeddings.
    embedding_model_name="databricks-bge-large-en" # The embedding model. Required for direct-access index or delta-sync index with self-managed embeddings.
)

CREATE OR REPLACE FUNCTION main.default.databricks_docs_vector_search (
  -- The agent uses this comment to determine how to generate the query string parameter.
  query STRING
  COMMENT 'The query string for searching Databricks documentation.'
) RETURNS TABLE
-- The agent uses this comment to determine when to call this tool. It describes the types of documents and information contained within the index.
COMMENT 'Executes a search on Databricks documentation to retrieve text documents most relevant to the input query.' RETURN
SELECT
  chunked_text as page_content,
  map('doc_uri', url, 'chunk_id', chunk_id) as metadata
FROM
  vector_search(
    -- Specify your Vector Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(
    function_names=[
        "main.default.databricks_docs_vector_search"
    ]
)
tools = toolkit.tools

向检索器工具添加跟踪

添加 MLflow 跟踪以监视和调试检索器。 通过跟踪可以查看每个执行步骤的输入、输出和元数据。

上一个示例将 @mlflow.trace 修饰器 添加到 __call__ 和分析方法。 修饰器创建一个 范围，该范围在调用函数时开始，并在函数返回时结束。 MLflow 会自动记录函数的输入和输出以及引发的任何异常。

import mlflow
from mlflow.entities import Document

# This code snippet has been truncated for brevity. See the full retriever example above.
class VectorSearchRetriever:
  ...

  # Create a RETRIEVER span. The span name must match the retriever schema name.
  @mlflow.trace(span_type="RETRIEVER", name="vector_search")
  def __call__(...) -> List[Document]:
    ...

  # Create a PARSER span.
  @mlflow.trace(span_type="PARSER")
  def parse_results(...) -> List[Document]:
    ...

若要验证代理评估和 AI Playground 等下游应用程序是否正确呈现检索器跟踪，请确保修饰器满足以下要求：

• 使用 MLflow 检索器范围架构 并验证函数是否返回 List[Document] 对象。
• 跟踪名称和 retriever_schema 名称必须匹配才能正确配置跟踪。 请参阅以下部分，了解如何设置检索器架构。

设置检索器架构以验证 MLflow 兼容性

如果从检索器返回的跟踪或 span_type="RETRIEVER" 不符合 MLflow 的标准检索器架构，则必须手动将返回的架构映射到 MLflow 的预期字段。 这会验证 MLflow 是否可以正确跟踪检索器并在下游应用程序中呈现跟踪。

手动设置检索器架构：

1. 定义代理时调用 mlflow.models.set_retriever_schema 。 用于 set_retriever_schema 将返回表中的列名映射到 MLflow 的预期字段，例如 primary_key， text_column和 doc_uri。

   # Define the retriever's schema by providing your column names
   mlflow.models.set_retriever_schema(
     name="vector_search",
     primary_key="chunk_id",
     text_column="text_column",
     doc_uri="doc_uri"
     # other_columns=["column1", "column2"],
   )
   

2. 通过提供 other_columns 字段的列名称列表，在检索器的架构中指定额外的列。

3. 如果有多个检索器，则可以为每个检索器架构使用唯一名称来定义多个架构。

代理创建期间设置的检索器架构会影响下游应用程序和工作流，例如评审应用和评估集。 具体而言，doc_uri 列充当检索器返回的文档的主要标识符。

• 审阅应用 显示 doc_uri，以帮助审阅者评估响应和跟踪文档来源。 请参阅评审应用 UI。
• 评估集 用于 doc_uri 将检索器结果与预定义评估数据集进行比较，以确定检索器的召回率和精度。 请参阅评估集（MLflow 2）。

从 Unity 目录卷读取文件

如果代理需要读取存储在 Unity 目录卷中的非结构化文件（文本文档、报表、配置文件等），则可以创建使用 Databricks SDK 文件 API 直接列出和读取文件的工具。

以下示例创建代理可以使用的两种工具：

• list_volume_files：列出卷中的文件和目录。
• read_volume_file：从卷中读取文本文件的内容。

LangChain/LangGraph

安装包含 Databricks AI Bridge 的 databricks-langchain 最新版本。

%pip install --upgrade databricks-langchain

from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>"  # TODO: Replace with your volume
w = WorkspaceClient()


@tool
def list_volume_files(directory: str = "") -> str:
    """Lists files and directories in the Unity Catalog volume.
    Provide a relative directory path, or leave empty to list the volume root."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    path = f"{base}/{directory.lstrip('/')}" if directory else base
    entries = []
    for f in w.files.list_directory_contents(path):
        kind = "dir" if f.is_directory else "file"
        size = f" ({f.file_size} bytes)" if not f.is_directory else ""
        entries.append(f"  [{kind}] {f.name}{size}")
    return "\n".join(entries) if entries else "No files found."


@tool
def read_volume_file(file_path: str) -> str:
    """Reads a text file from the Unity Catalog volume.
    Provide the path relative to the volume root, for example 'reports/q1_summary.txt'."""
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    full_path = f"{base}/{file_path.lstrip('/')}"
    resp = w.files.download(full_path)
    return resp.contents.read().decode("utf-8")

将工具绑定到 LLM 并运行工具调用循环：

from databricks_langchain import ChatDatabricks
from langchain_core.messages import HumanMessage, ToolMessage

llm = ChatDatabricks(endpoint="databricks-claude-sonnet-4-5")
llm_with_tools = llm.bind_tools([list_volume_files, read_volume_file])

messages = [HumanMessage(content="What files are in the volume? Can you read about_databricks.txt and summarize it in 2 sentences?")]
tool_map = {"list_volume_files": list_volume_files, "read_volume_file": read_volume_file}

for _ in range(5):  # max iterations
    response = llm_with_tools.invoke(messages)
    messages.append(response)
    if not response.tool_calls:
        break
    for tc in response.tool_calls:
        result = tool_map[tc["name"]].invoke(tc["args"])
        messages.append(ToolMessage(content=result, tool_call_id=tc["id"]))

print(response.content)

OpenAI

安装包含 Databricks AI Bridge 的 databricks-openai 最新版本。

%pip install --upgrade databricks-openai

from databricks.sdk import WorkspaceClient
from databricks_openai import DatabricksOpenAI
import json

VOLUME = "<catalog>.<schema>.<volume>"  # TODO: Replace with your volume
w = WorkspaceClient()
client = DatabricksOpenAI()

# Define the tool specifications
tools = [
    {
        "type": "function",
        "function": {
            "name": "list_volume_files",
            "description": "Lists files and directories in the Unity Catalog volume. Provide a relative directory path, or leave empty to list the volume root.",
            "parameters": {
                "type": "object",
                "properties": {
                    "directory": {
                        "type": "string",
                        "description": "Relative directory path within the volume. Leave empty for root.",
                    }
                },
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_volume_file",
            "description": "Reads a text file from the Unity Catalog volume. Provide the path relative to the volume root, for example 'reports/q4_summary.txt'.",
            "parameters": {
                "type": "object",
                "properties": {
                    "file_path": {
                        "type": "string",
                        "description": "Path to the file relative to the volume root.",
                    }
                },
                "required": ["file_path"],
            },
        },
    },
]


def execute_tool(name: str, args: dict) -> str:
    base = f"/Volumes/{VOLUME.replace('.', '/')}"
    if name == "list_volume_files":
        directory = args.get("directory", "")
        path = f"{base}/{directory.lstrip('/')}" if directory else base
        entries = []
        for f in w.files.list_directory_contents(path):
            kind = "dir" if f.is_directory else "file"
            size = f" ({f.file_size} bytes)" if not f.is_directory else ""
            entries.append(f"[{kind}] {f.name}{size}")
        return "\n".join(entries) if entries else "No files found."
    elif name == "read_volume_file":
        full_path = f"{base}/{args['file_path'].lstrip('/')}"
        resp = w.files.download(full_path)
        return resp.contents.read().decode("utf-8")
    return f"Unknown tool: {name}"


# Call the model with tools
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "List the files in the volume, then read about_databricks.txt and summarize it."},
]

response = client.chat.completions.create(
    model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
)

# Execute tool calls and send results back
while response.choices[0].finish_reason == "tool_calls":
    messages.append(response.choices[0].message)
    for tool_call in response.choices[0].message.tool_calls:
        args = json.loads(tool_call.function.arguments)
        result = execute_tool(tool_call.function.name, args)
        messages.append(
            {"role": "tool", "tool_call_id": tool_call.id, "content": result}
        )
    response = client.chat.completions.create(
        model="databricks-claude-sonnet-4-5", messages=messages, tools=tools
    )

print(response.choices[0].message.content)