Databricks Mosaic AI RAG Tutorial: Build a Production-Ready Knowledge Assistant

Overview

Retrieval-augmented generation (RAG) combines a language model with a searchable collection of trusted documents. Instead of asking a model to answer entirely from its training data, an application retrieves relevant passages at query time and places them in the model prompt. The model then uses that context to produce an answer, ideally with citations or links back to the source material.

A typical Mosaic AI RAG architecture contains these stages:

Source documents
      ↓
Ingestion and parsing
      ↓
Chunking and metadata enrichment
      ↓
Embeddings and indexed storage
      ↓
Databricks Vector Search retrieval
      ↓
Prompt construction with retrieved context
      ↓
Model endpoint and response
      ↓
Evaluation, logging, feedback, and monitoring

Databricks can provide a useful foundation because the documents, transformations, evaluation data, and application code can be managed within a common data and AI workflow. The right design still depends on the use case. A small internal policy assistant may need only a modest document set and a carefully defined access boundary. A customer-facing assistant may require stricter authorization, higher availability, response filtering, and more extensive evaluation.

Before implementation, write down the assistant's intended questions, allowed sources, users, response format, and failure behavior. “Answer questions about company documents” is too broad to evaluate. “Answer questions about the current employee travel policy, cite the policy section, and say when the documents do not contain an answer” is a more testable objective.

Checklist by scenario

Scenario 1: Preparing a document collection

• Define source ownership. Record who owns each document set, how often it changes, and which version should be authoritative.
• Choose supported formats deliberately. PDFs, HTML pages, office files, tickets, and database records often require different parsing and cleanup steps.
• Preserve useful structure. Keep titles, headings, section names, publication dates, document identifiers, and URLs as metadata where possible.
• Remove or isolate unwanted content. Navigation menus, repeated footers, boilerplate, stale drafts, and duplicated pages can reduce retrieval quality.
• Define an update path. Store raw inputs separately from normalized text so that parsing logic can be revised without losing the original material.
• Record access attributes. A document's department, sensitivity, tenant, region, or user group may be needed for metadata filtering and authorization.

Use Delta tables or another governed storage layer as the durable system of record for processed content. The vector index should be treated as a retrieval layer that can be rebuilt or synchronized from that record. For table maintenance and file layout considerations, see the Delta Lake Maintenance Guide.

Scenario 2: Chunking and indexing content

• Start with semantic boundaries. Prefer sections, paragraphs, list groups, or policy clauses over arbitrary cuts through a sentence.
• Keep chunks focused. A chunk should contain enough context to answer a question but not so many unrelated topics that similarity search becomes ambiguous.
• Use overlap selectively. Overlap can preserve context across boundaries, but excessive overlap increases duplication and may crowd the retrieved context.
• Attach stable identifiers. Every chunk should be traceable to its source document, section, version, and ingestion run.
• Plan for re-indexing. Decide how changed and deleted documents will be detected, updated, and removed from the search index.
• Test retrieval independently. Before connecting an LLM, check whether relevant chunks appear for representative questions.

Databricks Vector Search can support the retrieval stage, but index configuration is not a substitute for clean source data. Compare different chunking approaches with a fixed question set rather than selecting a strategy based on a single successful demonstration. The Databricks Vector Search Guide is useful background when planning index setup and operational tradeoffs.

Scenario 3: Designing the RAG prompt

• State the assistant's role and scope. Explain what kind of questions it can answer and which requests are outside scope.
• Separate instructions from retrieved text. Use clear delimiters so document content is treated as evidence rather than as a new instruction.
• Require grounded behavior. Tell the model to use the supplied context, distinguish facts from uncertainty, and decline when the evidence is insufficient.
• Specify the response format. Define whether the answer should include citations, a short summary, numbered steps, or a structured JSON response.
• Handle conversation history carefully. Include only the history needed to resolve the current question and avoid allowing old assumptions to override current retrieved evidence.
• Make citations traceable. Pass source labels or document links with each chunk so the final answer can identify where claims came from.

A practical prompt pattern is: answer the user's question using only the retrieved context; cite the relevant source identifiers; explain uncertainty; and say that the available documents do not establish an answer when they do not. This will not eliminate hallucinations, but it gives evaluation and monitoring a clear behavioral target.

Scenario 4: Evaluating and deploying the application

• Build a representative evaluation set. Include common questions, ambiguous questions, questions with no answer, outdated-document cases, and permission-sensitive requests.
• Measure retrieval and generation separately. A fluent answer may still be based on the wrong passage. Review retrieved relevance before judging the generated response.
• Include human review. Automated checks can score format, citation presence, or similarity, but subject-matter experts are needed for factual and policy judgments.
• Test failure modes. Try misspellings, follow-up questions, conflicting documents, prompt injection attempts, and requests for restricted information.
• Version the important components. Track source snapshots, chunking logic, embedding configuration, retrieval settings, prompt templates, model selection, and evaluation results.
• Choose a serving path with operational ownership. Define authentication, scaling expectations, timeouts, retries, logging, and rollback procedures before exposing the application.

For model endpoint lifecycle considerations, connect this work with the Databricks Model Serving Guide. If ingestion and evaluation run on a schedule, use explicit dependencies and failure handling; the Databricks Jobs Guide covers the surrounding workflow concerns.

What to double-check

Retrieval quality

Ask whether the correct document is in the corpus, whether the parser preserved the relevant text, and whether the query and passage use compatible terminology. If the right passage is not retrieved, changing the model prompt alone is unlikely to solve the problem. Test metadata filters, top-result counts, query rewriting, and hybrid retrieval where keyword matches matter as much as semantic similarity.

Authorization and data boundaries

Do not assume that hiding a source link protects restricted content. Enforce access at the data and application layers, and ensure that retrieved chunks are filtered for the requesting user's permissions before they reach the model. Review service principals, secrets, network exposure, audit events, and workspace permissions. Use the Databricks Security Best Practices Checklist as a companion review.

Freshness and conflicting sources

Define what “current” means for the use case. A knowledge assistant should not silently combine an old procedure with a newer one. Store effective dates or versions, prioritize authoritative sources, and decide how the application should respond when documents disagree. Test the update process by changing a source document and confirming that the expected content becomes retrievable.

Observability and cost control

Log enough information to investigate failures without storing sensitive user content unnecessarily. Useful fields can include request identifiers, retrieval scores, source identifiers, latency by stage, token or payload sizes, model response status, and user feedback. Establish practical limits for query length, retrieved context, retries, and concurrent traffic. Review usage patterns before increasing context or model size.

Common mistakes

• Starting with the model instead of the question. A stronger model cannot compensate for an undefined scope or poor source coverage.
• Indexing raw documents without inspection. Broken extraction, duplicated pages, and missing headings often create retrieval failures that are misdiagnosed as model problems.
• Using one chunking rule for every source. A legal clause, product table, and support ticket have different structural needs.
• Putting the entire document set into the prompt. More context is not automatically better; irrelevant passages can dilute evidence and increase response cost or latency.
• Evaluating only happy-path questions. Include unanswerable, adversarial, ambiguous, and permission-sensitive examples.
• Allowing retrieved text to override application rules. Retrieved content is data, not trusted instructions. Delimit it and apply policy controls outside the model.
• Skipping source citations. Traceability helps users verify answers and helps developers identify which document or chunk caused an error.
• Deploying without a rollback plan. Keep prior index, prompt, and application versions available so a bad ingestion run or configuration change can be reversed.

When to revisit

Revisit this checklist before seasonal planning cycles, major content refreshes, or any change to the assistant's audience and scope. Also review it whenever workflows or tools change, including a new embedding model, model endpoint, parser, index configuration, authentication method, or deployment process.

At minimum, schedule a recurring review of four signals: unanswered user questions, incorrect or weakly supported answers, stale or unauthorized retrievals, and changes in latency or usage. Sample production conversations according to your privacy and retention requirements, then add representative failures to the evaluation set. A RAG application improves when each incident becomes a test case rather than an isolated fix.

Use this final action list for the next iteration:

1. Write five to ten target questions and define what a good answer must contain.
2. Trace each expected answer to an approved source document and section.
3. Inspect parsed text and metadata before creating or rebuilding the vector index.
4. Test retrieval independently, then test the full prompt-and-model response.
5. Run authorization, stale-content, prompt-injection, and no-answer tests.
6. Record versions and evaluation results before deployment.
7. After release, monitor failures and update the test set whenever the source corpus or workflow changes.

This process keeps a Databricks knowledge assistant maintainable. The durable asset is not only the endpoint or prompt; it is the governed source pipeline, measurable retrieval behavior, clear failure policy, and repeatable deployment checklist around them.