Skip to main content
Last updated on

Connect agents to unstructured data

Agents often need to query unstructured data like document collections, knowledge bases, or text corpora to answer questions and provide context-aware responses.

For new Python agents, use AI Bridge (recommended) to query AI Search indexes directly. For existing MCP integrations, see the legacy workspace endpoint. You can also connect to external vector stores or read files from Unity Catalog volumes.

For new Python agents that query Databricks AI Search indexes, Databricks recommends AI Bridge. Use VectorSearchRetrieverTool from databricks-langchain or databricks-openai to expose an index as an agent tool.

Before running these examples, sign in to your workspace and choose an AI Search index that you can query. Replace the example index name with its three-part name. The first examples use Databricks managed embeddings and a model endpoint that supports tool calling.

Run the installation commands in a terminal. In a notebook, use %pip instead of pip and restart Python after installation.

Install the latest version of databricks-langchain and LangGraph:

Bash
pip install --upgrade databricks-langchain langgraph

Create a retriever tool and give it a description that tells the agent when to search. This example tests retrieval directly, then runs an agent that can call the tool and use its results.

Python
from databricks_langchain import VectorSearchRetrieverTool, ChatDatabricks
from langgraph.prebuilt import create_react_agent

vs_tool = VectorSearchRetrieverTool(
    index_name="catalog.schema.my_databricks_docs_index",
    tool_name="databricks_docs_retriever",
    tool_description="Search Databricks documentation to answer product questions.",
)
print(vs_tool.invoke("What is Databricks Agent Framework?"))

agent = create_react_agent(
    ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
    tools=[vs_tool],
)
result = agent.invoke({"messages": [{
    "role": "user",
    "content": "Search the documentation and explain Databricks Agent Framework.",
}]})
for message in result["messages"]:
    print(message)

Check the printed conversation for a tool call, retrieved documents, and an answer based on those documents.

For scenarios that use either direct-access indexes or Delta Sync indexes using self-managed embeddings, you must configure the VectorSearchRetrieverTool and specify a custom embedding model and text column. See options for providing embeddings.

The following example shows you how to configure a VectorSearchRetrieverTool with columns and embedding keys.

Python
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="databricks_docs_retriever", # 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.
)

For additional details, see the API docs for VectorSearchRetrieverTool.

After testing the retriever, add it to your custom agent.

Query a vector search index outside Databricks​

If your vector index is hosted outside Databricks, you can create a Unity Catalog connection to connect to the external service and use the connection in your agent code. See Connect to external HTTP services.

The following example queries an Azure AI Search index through an HTTP connection. You need a connection with permission to query that index and USE CONNECTION on the connection. Install requests, databricks-sdk, and mlflow before running the Python example.

  1. Create a Unity Catalog connection to Azure AI Search using HTTP connection setup. This example uses a bearer token; choose the authentication method your Azure AI Search deployment requires.

    SQL
    CREATE CONNECTION azure_search_connection
    TYPE HTTP
    OPTIONS (
      host 'https://example.search.windows.net',
      base_path '/',
      bearer_token secret ('<secret-scope>','<secret-key>')
    );

  2. Define the retriever tool in agent code using the Unity Catalog connection. This example uses MLflow decorators to enable agent tracing.

    note

    To conform to the MLflow retriever schema, the retriever function should return a List[Document] object and use the metadata field in the Document class to add additional attributes to the returned document, such as doc_uri and similarity_score. See MLflow Document.

    Python
    import mlflow
    import requests
    from databricks.sdk import WorkspaceClient
    from mlflow.entities import Document

    class VectorSearchRetriever:
        def __init__(self, connection_name, index_name):
            self.workspace = WorkspaceClient()
            self.connection_name = connection_name
            self.index_name = index_name

        @mlflow.trace(span_type="RETRIEVER", name="vector_search")
        def __call__(self, query_vector: list[float], score_threshold: float = 0.0) -> list[Document]:
            response = requests.post(
                f"{self.workspace.config.host}/api/2.0/unity-catalog/connections/"
                f"{self.connection_name}/proxy/indexes/{self.index_name}/docs/search",
                params={"api-version": "2024-07-01"},
                headers=self.workspace.config.authenticate(),
                json={
                    "select": "HotelId,HotelName,Description,Category",
                    "vectorQueries": [{
                        "vector": query_vector,
                        "k": 7,
                        "fields": "DescriptionVector",
                        "kind": "vector",
                        "exhaustive": True,
                    }],
                },
                timeout=30,
            )
            response.raise_for_status()
            return self.parse_results(response.json(), score_threshold)

        @mlflow.trace(span_type="PARSER")
        def parse_results(self, results, score_threshold) -> list[Document]:
            return [
                Document(
                    page_content=item.get("Description", ""),
                    id=item.get("HotelId"),
                    metadata={
                        "similarity_score": item.get("@search.score", 0),
                        "HotelName": item.get("HotelName"),
                        "Category": item.get("Category"),
                    },
                )
                for item in results.get("value", [])
                if item.get("@search.score", 0) >= score_threshold
            ]

  3. Generate a query embedding with the same model and dimensions used by your index. Save the numeric array as query-vector.json, then run the retriever:

    Python
    import json
    from pathlib import Path

    query_vector = json.loads(Path("query-vector.json").read_text())
    retriever = VectorSearchRetriever("azure_search_connection", "hotels_vector_index")
    results = retriever(query_vector, score_threshold=0.1)
    for document in results:
        print(document.page_content)

Add tracing to a retriever tool​

Add MLflow tracing to monitor and debug your retriever. Tracing lets you view inputs, outputs, and metadata for each step of execution.

The previous example adds the @mlflow.trace decorator to both the __call__ and parsing methods. The decorator creates a span that starts when the function is invoked and ends when it returns. MLflow automatically records the function's input and output and any exceptions raised.

note

LangChain, LlamaIndex, and OpenAI library users can use MLflow auto logging in addition to manually defining traces with the decorator. See Tracing overview.

Return a list of Document objects with page_content and metadata, as the example above does. This follows the MLflow retriever span schema, so evaluation and review tools can recognize the retrieved documents.

Read files from a Unity Catalog volume​

If your agent needs to read unstructured files (text documents, reports, configuration files, etc.) stored in a Unity Catalog volume, you can create tools that use the Databricks SDK Files API to list and read files directly.

Choose an existing volume and grant your agent's identity USE CATALOG, USE SCHEMA, and READ VOLUME. Replace <catalog>.<schema>.<volume> in the code. For the sample task, use a UTF-8 text file named about_databricks.txt, or change the prompt to a file in your volume.

The following examples create two tools your agent can use:

  • list_volume_files: Lists files and directories in the volume.
  • read_volume_file: Reads the contents of a text file from the volume.

Install the latest version of databricks-langchain that includes Databricks AI Bridge.

Bash
pip install --upgrade databricks-langchain langgraph
Python
from pathlib import PurePosixPath
from databricks.sdk import WorkspaceClient
from langchain_core.tools import tool

VOLUME = "<catalog>.<schema>.<volume>"
w = WorkspaceClient()


def volume_path(relative_path: str = "") -> str:
    path = PurePosixPath(relative_path)
    if path.is_absolute() or ".." in path.parts:
        raise ValueError("Use a path relative to the configured volume.")
    return str(PurePosixPath("/Volumes", *VOLUME.split(".")) / path)


@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."""
    path = volume_path(directory)
    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'."""
    full_path = volume_path(file_path)
    resp = w.files.download(full_path)
    return resp.contents.read().decode("utf-8")

Bind the tools to an LLM and run a tool-calling loop:

Python
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"]))

else:
    raise RuntimeError("The agent reached the tool-calling limit.")

print(response.content)

Query an AI Search index using MCP (legacy)​

The AI Search MCP server uses a legacy workspace MCP endpoint. It requires an index with Databricks managed embeddings. See Create AI Search endpoints and indexes.

The legacy workspace MCP URL for AI Search is https://<workspace-hostname>/api/2.0/mcp/ai-search/{catalog}/{schema}/{index_name}. Connect to it and list the tools it exposes:

Python
from databricks.sdk import WorkspaceClient
from databricks_mcp import DatabricksMCPClient

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

mcp_client = DatabricksMCPClient(
    server_url=f"{host}/api/2.0/mcp/ai-search/<catalog>/<schema>/<index-name>",
    workspace_client=workspace_client,
)
tools = mcp_client.list_tools()

To build and deploy an agent that uses this server, see Use MCP tools in a Python agent. Grant the agent SELECT on the index's Unity Catalog securable.

AI Search was formerly Vector Search. The previous /api/2.0/mcp/vector-search/ URL prefix still works.

Query Databricks AI Search using UC functions (deprecated)​

note

For new Python agents, use AI Bridge. The following example is a reference for existing integrations.

You can create a Unity Catalog function that wraps an AI Search index query. This approach:

  • Supports production use cases with governance and discoverability
  • Uses the vector_search() SQL function under the hood
  • Supports automatic MLflow tracing
    • You must align the function's output to the MLflow retriever schema by using the page_content and metadata aliases.
    • Any additional metadata columns must be added to the metadata column using the SQL map function, rather than as top-level output keys.

Run the following code in a notebook or SQL editor to create the function:

SQL
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 AI Search index name here
    index => 'catalog.schema.databricks_docs_index',
    query => query,
    num_results => 5
  )

Follow the function tool setup, then wrap this retriever with UCFunctionToolkit to use it in your agent. This enables automatic tracing through MLflow by automatically generating RETRIEVER span types in MLflow logs.

Python
from databricks_langchain import UCFunctionToolkit
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

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

Unity Catalog retriever tools have the following caveats:

  • SQL clients might limit the maximum number of rows or bytes returned. To prevent data truncation, truncate column values returned by the UDF. For example, you could use substring(chunked_text, 0, 8192) to reduce the size of large content columns and avoid row truncation during execution.
  • Because this tool is a wrapper for the vector_search() function, it is subject to the same limitations as the vector_search() function. See Limitations.

For more information about UCFunctionToolkit see Unity Catalog documentation.

Custom retriever schemas (legacy)​

mlflow.models.set_retriever_schema is deprecated. For new retrievers, return Document objects using the standard retriever schema. The following configuration is for existing integrations that return custom fields.

To set the retriever schema manually:

  1. Call mlflow.models.set_retriever_schema when you define your agent. Use set_retriever_schema to map the column names in the returned table to MLflow's expected fields such as primary_key, text_column, and doc_uri.

    Python
    import mlflow

    # 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. Specify additional columns in your retriever's schema by providing a list of column names with the other_columns field.

  3. If you have multiple retrievers, you can define multiple schemas by using unique names for each retriever schema.

The retriever schema set during agent creation affects downstream applications and workflows, such as the review app and evaluation sets. Specifically, the doc_uri column serves as the primary identifier for documents returned by the retriever.

  • The review app displays the doc_uri to help reviewers assess responses and trace document origins. See Review App UI.
  • Evaluation sets use doc_uri to compare retriever results against predefined evaluation datasets to determine the retriever's recall and precision. See Evaluation dataset reference.
On this page