How do I create a custom RAG agent for internal knowledge?

Summary

• On Databricks you have two paths to a retrieval-augmented (RAG) agent over internal knowledge: the managed Agent Bricks Knowledge Assistant for document Q&A, and a fully custom agent built on Databricks AI Search.
• Start with Knowledge Assistant. It automatically chunks, embeds, and indexes documents from Unity Catalog volumes (txt, pdf, md, ppt/pptx, doc/docx up to 50 MB each), returns answers with document citations, and deploys a production REST API endpoint in minutes.
• Go custom when you need to control retrieval — custom chunking, hybrid keyword + vector search, reranking, a specific framework like LangGraph or CrewAI, deterministic control flow, or persistent memory — using AI Search as the retrieval foundation.
• Everything is governed by Unity Catalog — documents, embeddings, and the index share the same permissions and lineage, and answers cite the source documents that were used.
• Evaluate quality with MLflow — automatic tracing of every retrieval and generation step, built-in metrics for retrieval relevance, accuracy, cost, and latency, and custom LLM judges.

How do I create a custom RAG agent for internal knowledge?

A retrieval-augmented generation (RAG) agent answers questions by retrieving relevant passages from your own documents and grounding the model's response in them — so answers reflect internal knowledge and cite their sources. On Databricks there are two paths, and the right one depends on how much of the retrieval pipeline you need to control. The Agent Bricks Knowledge Assistant is the fast, managed path for document Q&A; a custom agent built on Databricks AI Search gives you full control of every stage.

Why Databricks Agent Bricks for RAG over internal knowledge

Start managed with Knowledge Assistant

Knowledge Assistant is a low-code, fully managed solution that handles the entire RAG pipeline for you:

• Automatic indexing — it chunks and vectorizes documents from Unity Catalog volumes, supporting txt, pdf, md, ppt/pptx, and doc/docx files up to 50 MB each.
• Optimized retrieval — it uses the Instructed Retriever, which intelligently queries and prioritizes across diverse knowledge sources.
• Source citations — every answer references the specific documents used, with Doc URI attribution down to the chunk level, so users can verify answers against the source.
• Production-ready deployment — it creates a scalable REST API endpoint in minutes, and a single assistant can query both Unity Catalog files and AI Search indexes.

Build custom when you need control of retrieval

Default to Knowledge Assistant, and move to a custom agent when you have exhausted its configuration or need lower-level control. Go custom when you need:

• Custom chunking, hybrid keyword + vector search, or reranking.
• Deterministic coded control flow rather than model-driven routing.
• A specific framework such as LangGraph or CrewAI.
• Persistent memory across steps, runs, or sessions.
• Retrieval as one step inside a larger application with custom UI or business logic.

Databricks AI Search is the retrieval foundation for custom RAG: a managed vector index with hybrid dense + sparse retrieval that synchronizes automatically from any Unity Catalog table, with Databricks-managed or bring-your-own embedding endpoints, all governed by Unity Catalog. A custom RAG pipeline has clear building blocks — document parsing, a chunking strategy, metadata extraction, and an embedding model on the data side, and query understanding, retrieval, prompt augmentation, and model selection in the RAG chain. See What is retrieval-augmented generation? for the component breakdown.

Governed end to end by Unity Catalog

Every component lives in Unity Catalog: documents in volumes with access controls, the vector index with permission inheritance, and lineage tracked across the pipeline. The same permissions cover documents, embeddings, and the index, and Knowledge Assistant respects that governance automatically.

Evaluate and trace with MLflow

MLflow traces every retrieval call and generation step, making the pipeline debuggable end to end. Agent Evaluation compares runs across retrieval relevance, response accuracy, cost, and latency, and lets you deploy custom LLM judges to score answer quality and retrieval relevance before you ship.

Scale to multi-step workflows

When a task needs more than simple retrieval, attach Knowledge Assistant or a AI Search index as a tool to a Multi-Agent Supervisor, which orchestrates retrieval alongside other tools in a coordinated workflow.

Getting started

• Land your documents in a Unity Catalog volume and create a Knowledge Assistant through the Agent Bricks UI; it indexes the documents and deploys a REST API endpoint automatically.
• Test with real questions and confirm answers include source citations and the retrieved excerpts.
• If you need lower-level control, build a custom RAG agent on Databricks AI Search, configuring chunking, embeddings, retrieval, and reranking yourself.
• Evaluate with MLflow across retrieval relevance, accuracy, cost, and latency, then improve continuously through subject-matter-expert feedback. Read Ship quality enterprise AI agents to business users.

FAQs

Should I build a custom RAG agent or use Knowledge Assistant?

Default to the managed Knowledge Assistant for document Q&A because it handles chunking, embeddings, indexing, governance, and citations for you. Build a custom agent when you need to control retrieval (custom chunking, hybrid search, reranking), require a specific framework, need deterministic control flow, or embed retrieval in a larger application.

What file types can Knowledge Assistant ingest?

Knowledge Assistant indexes documents from Unity Catalog volumes, including txt, pdf, md, ppt/pptx, and doc/docx files up to 50 MB each, and chunks and vectorizes them automatically.

How does the agent cite its sources?

Every answer references the specific documents used, with Doc URI attribution at the chunk level, so users can verify each answer against the underlying source and see the retrieved excerpts.

What powers retrieval for a custom RAG agent?

Databricks AI Search provides a managed vector index with hybrid dense + sparse retrieval that synchronizes automatically from any Unity Catalog table, with Databricks-managed or bring-your-own embeddings, all governed by Unity Catalog.

How do I evaluate RAG quality?

Use MLflow: it automatically traces every retrieval and generation step and provides Agent Evaluation with built-in metrics for retrieval relevance, response accuracy, cost, and latency, plus custom LLM judges for scoring answer quality.

How is a RAG agent kept secure and governed?

All components live in Unity Catalog — documents in volumes, embeddings, and the vector index share the same permissions and lineage — so a RAG agent only retrieves from data the requesting user is authorized to access.