RAG Retrieval for Agents: Querying AI Search and Injecting Chunks into Context

1.The RAG chain at request time

Once the index exists, every user request goes through the RAG chain. An optional preprocessing step can rewrite the request, put it into a template, or pull out keywords, producing a *retrieval query*. During retrieval, that query is embedded with the same model used for the chunks, compared against them (for example by cosine similarity), and the top-ranked results come back. Prompt augmentation builds the LLM prompt by putting the user's query and the retrieved context into a template that tells the model how to use each part. In generation, the LLM answers from that context. An optional post-processing step can add citations or apply business rules. Guardrails can apply at any step, including checking user permissions before the chain touches a data source.

Sources1

2.Give the agent a retriever over the index

An agent reaches the index through a tool. Databricks documents three ways to set one up, depending on where the index lives and how much control you need.

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

VectorSearchRetrieverTool lets you shape retrieval in code. You set how many chunks come back (num_results), which columns are returned, which filters apply, and whether the query is ANN or HYBRID. If the index is a Direct Vector Access Index, or a Delta Sync Index with self-managed embeddings, the tool cannot embed the query by itself. You must pass it a text_column and an embedding model.

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

Sources2

3.How retrieved chunks get into the agent's context

When retrieval is exposed as a tool, the LLM decides when to call it. Its description is the main signal for that decision, so write a specific tool_description. After the model calls the tool, your code runs the search and appends the results to the conversation as a tool message linked to the call ID. The model's next turn then sees the retrieved chunks next to the user's question. This is the augmentation step in practice.

messages.append(first_response.choices[0].message)
messages.append({
  "role": "tool",
  "tool_call_id": tool_call.id,
  "content": json.dumps(result)
})

If you write your own retriever, follow the MLflow retriever schema: return a List[Document] and put attributes such as doc_uri and similarity_score in each Document's metadata field. Tracing can then record what was retrieved, and the generation step has what it needs to cite its sources.

Sources12

4.Unity Catalog decides what the retriever can reach

An AI Search index is a Unity Catalog securable, so querying it follows the usual privilege chain. A user or principal that does not own the index needs USE CATALOG on its catalog, USE SCHEMA on its schema, and SELECT on the index. For the managed MCP server, the documentation says to grant the agent SELECT on the index's securable. This is how governance carries through to retrieval: the corpus lives in a governed table, the index is a governed object, and the agent can query only what it has been granted.

Sources23

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.Granting SELECT on the AI Search index is enough for a non-owner to query it.Why is that wrong?

   A non-owner also needs USE CATALOG on the catalog and USE SCHEMA on the schema that contain the index.

   Covered in Unity Catalog decides what the retriever can reach

2. 2.VectorSearchRetrieverTool needs only index_name, whatever kind of index it points at.Why is that wrong?

   For a Direct Vector Access Index, or a Delta Sync Index with self-managed embeddings, you must also give it a text column and an embedding model so it can embed the query.

   Covered in Give the agent a retriever over the index

Practise it for real

Prototype a retriever tool over an AI Search index and confirm that an LLM calls it and uses the retrieved chunks.

1. 1.Run %pip install --upgrade databricks-langchain in a notebook.

   Why: This version includes Databricks AI Bridge, which provides VectorSearchRetrieverTool.

   You should see: The package installs and you can import from databricks_langchain.

2. 2.Create VectorSearchRetrieverTool(index_name="catalog.schema.my_databricks_docs_index", tool_name=..., tool_description=...) with your own index name and a specific description.

   Why: The LLM uses tool_name and tool_description to decide when to call the retriever.

   You should see: The tool object is created. If it fails with a permissions error, check USE CATALOG, USE SCHEMA and SELECT on the index.

3. 3.Call vs_tool.invoke("<a question your corpus answers>").

   Why: This checks retrieval by itself before an LLM is involved.

   You should see: Relevant chunks from your index come back.

4. 4.Create ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"), call bind_tools([vs_tool]) on it, then invoke it with a question that starts "Based on the documentation…".

   Why: This tests whether the model chooses to call the retriever and answers from the retrieved context.

   You should see: The model makes a tool call to your retriever and answers from what the index returned.