Building RAG Pipelines with Databricks

Table of Contents

1. Course Overview
2. Module 1 — Getting Started with RAG on Databricks
   • 2.1 Introduction to Databricks
   • 2.2 Understanding the RAG Pattern
   • 2.3 Scenario and Environment Setup
3. Module 2 — Embeddings and Mosaic AI Vector Search
   • 3.1 Understanding Embeddings and Text Similarity
   • 3.2 Preparing, Chunking, and Storing Data in a Delta Table
   • 3.3 Setting Up Mosaic AI Vector Search
   • 3.4 Indexing and Querying Data in Vector Search
4. Module 3 — Building the RAG Pipeline with Databricks
   • 4.1 Understanding the RAG Workflow in Databricks
   • 4.2 Prompt Engineering for RAG
   • 4.3 Building the RAG Pipeline
   • 4.4 Tracking and Evaluating the RAG Pipeline with MLflow
   • 4.5 Registering and Deploying the RAG Model
5. Module 4 — Agent Bricks for RAG Workflows
   • 5.1 Understanding Agent Bricks in Databricks Mosaic AI
   • 5.2 Building an Agentic RAG with Agent Bricks
6. Environment Setup
7. Code File Reference

1. Course Overview

┌─────────────────────────────────────────────────────────────┐
│           COURSE OUTLINE                                     │
├─────────────────────────────────────────────────────────────┤
│  Module 1 │ Getting Started with RAG on Databricks          │
│  Module 2 │ Embeddings and Mosaic AI Vector Search          │
│  Module 3 │ Building the RAG Pipeline                       │
│  Module 4 │ Agent Bricks for RAG Workflows                  │
└─────────────────────────────────────────────────────────────┘

This course covers building end-to-end RAG (Retrieval-Augmented Generation) pipelines on Databricks, from understanding the foundational concepts to deploying autonomous agents.

2. Module 1 — Getting Started with RAG on Databricks

2.1 Introduction to Databricks

Databricks is a cloud-based data intelligence platform. It unifies different use cases on a single platform:

┌──────────────────────────────────────────────────────────────────────┐
│                    DATABRICKS LAKEHOUSE ARCHITECTURE                  │
├──────────────────────────────────────────────────────────────────────┤
│                                                                       │
│  ┌──────────────┐  ┌──────────────┐  ┌──────────────────────────┐   │
│  │  Data Eng.   │  │  Mosaic AI   │  │    Databricks SQL        │   │
│  │  (ETL/ELT)   │  │   / ML       │  │   (Data Warehousing)     │   │
│  └──────────────┘  └──────────────┘  └──────────────────────────┘   │
│                                                                       │
│  ┌──────────────────────────────────────────────────────────────┐    │
│  │              COMPUTE ENGINES (Cluster / Serverless)          │    │
│  │                    Apache Spark (open-source)                │    │
│  └──────────────────────────────────────────────────────────────┘    │
│                                                                       │
│  ┌──────────────────────────────────────────────────────────────┐    │
│  │         UNITY CATALOG — Governance, cataloging, security     │    │
│  └──────────────────────────────────────────────────────────────┘    │
│                                                                       │
│  ┌──────────────────────────────────────────────────────────────┐    │
│  │          DELTA LAKE — ACID Storage Layer                     │    │
│  └──────────────────────────────────────────────────────────────┘    │
│                                                                       │
│  ┌───────────┐     ┌─────────────────┐     ┌──────────────────┐     │
│  │    AWS    │     │  Microsoft Azure │     │  Google Cloud    │     │
│  └───────────┘     └─────────────────┘     └──────────────────┘     │
└──────────────────────────────────────────────────────────────────────┘

Key components:

2.2 Understanding the RAG Pattern

Challenges of LLMs without RAG

graph LR
    U[👤 User] -->|Query| M[GenAI Model]
    M -->|Generic response| U

    style M fill:#f9a825,color:#000

Problems with classic LLMs:

• 🔴 Hallucinations — incorrect or fabricated responses
• 🔴 No access to enterprise data — trained on public data only
• 🔴 No real-time data access — fixed training cutoff
• 🔴 Expensive retraining — significant time and cost
• 🔴 Limited context window — restricted number of tokens

The solution: the RAG pattern

RAG is a pattern that generates responses via an LLM, grounded in external knowledge sources.

Concrete example — Health insurance policy:

┌─────────────────────────────────────────────────────────────────────┐
│  WITHOUT RAG                                                         │
│                                                                      │
│  Question: "How many days are covered before hospitalization?"       │
│  Answer: "Generally, between 30 and 60 days. Please consult your    │
│           policy." (generic response)                                │
├─────────────────────────────────────────────────────────────────────┤
│  WITH RAG (on the user's insurance policy)                          │
│                                                                      │
│  Question: "How many days are covered before hospitalization?"       │
│  Answer: "According to your policy, coverage is 90 days."           │
│           (precise response grounded in data)                        │
└─────────────────────────────────────────────────────────────────────┘

The two phases of RAG

flowchart LR
    subgraph PHASE_A["Phase A — Data Preparation"]
        D[Enterprise Documents] --> TS[Text Splitter]
        TS --> C[Chunks]
        C --> SE[Search Engine / Index]
    end

    subgraph PHASE_B["Phase B — Query Response"]
        U[👤 User] -->|Question| SE2[Search Engine]
        SE2 -->|Similarity Search| RC[Relevant Chunks]
        RC --> LLM[LLM]
        U -->|Question| LLM
        LLM -->|Natural language response| U
    end

    style PHASE_A fill:#e3f2fd
    style PHASE_B fill:#e8f5e9

Detailed steps:

Phase A — Data Preparation:
──────────────────────────────────
Documents → [Text Splitter] → Chunks → [Embedding Model] → Vectors → [Vector Search Index]

Phase B — Query Response:
─────────────────────────────────
                    ┌─────────────────────────────────┐
                    │  1. RETRIEVAL                   │
User question → [Embedding Model] → Question Vector
Question Vector → [Vector Search] → Similar Chunks

                    │  2. AUGMENTATION                │
Question + Relevant Chunks → Enriched Prompt

                    │  3. GENERATION                  │
Enriched Prompt → [LLM] → Response grounded in data
                    └─────────────────────────────────┘

Advantages of RAG

Use Cases

┌─────────────────────────────────────────────────────────────────────┐
│                        RAG USE CASES                                │
├───────────────────┬─────────────────────────────────────────────────┤
│ 🏢 Enterprise     │ Knowledge assistants, IT helpdesk,              │
│                   │ HR policies                                      │
├───────────────────┼─────────────────────────────────────────────────┤
│ 🎧 Customer Support│ FAQ, product documentation                     │
├───────────────────┼─────────────────────────────────────────────────┤
│ 🏥 Healthcare     │ Clinical guidelines, medical research           │
├───────────────────┼─────────────────────────────────────────────────┤
│ 💰 Finance        │ Risk analysis, investment research               │
├───────────────────┼─────────────────────────────────────────────────┤
│ ⚖️ Legal          │ Regulatory compliance, legal preparation        │
└───────────────────┴─────────────────────────────────────────────────┘

2.3 Scenario and Environment Setup

Course Scenario

The scenario used throughout the course is a health insurance policy (52-page document). A typical policy includes:

• Definitions
• Covered benefits
• Exclusions
• General conditions
• Appendices

Environment Setup

Prerequisites:

• Databricks Free Edition OR Azure Databricks with Unity Catalog configured
• Serverless compute

Setup steps:

flowchart TD
    A[Create an 'insurance' Catalog in Unity Catalog] --> B[Create a 'rag' Schema]
    B --> C[Create a 'raw' Volume]
    C --> D[Upload the PDF policy to the Volume]
    D --> E[✅ Environment ready]

In the Databricks interface:

1. Go to Catalog → Settings → Verify Unity Catalog Metastore
2. Create a new catalog: insurance
3. Create a schema: rag
4. Create a volume: raw
5. Upload the Contoso Premier Policy.pdf file to the volume

3. Module 2 — Embeddings and Mosaic AI Vector Search

3.1 Understanding Embeddings and Text Similarity

Limitations of Keyword Search

Text A: "Xbox is an excellent gaming console"
         Keywords: Xbox, excellent, gaming, console

Text B: "I can spend all day with PlayStation"
         Keywords: spend, day, PlayStation

→ No common keywords!
→ Yet, both sentences talk about gaming consoles.
→ Keyword search MISSES this relationship.

Concept of Embeddings

An embedding is a vector of floating-point numbers that represents the semantic meaning of text in a multidimensional space.

                    ┌─────────────────────┐
"Xbox is an         │   Embeddings Model  │ → [ 0.17, -1.75, ..., 7.28 ]
 excellent gaming   │                     │   (Embedding for Text A)
 console"           └─────────────────────┘

"I can spend        ┌─────────────────────┐
 all day with       │   Embeddings Model  │ → [ 3.29,  1.18, ..., -0.28 ]
 PlayStation"       └─────────────────────┘   (Embedding for Text B)

Cosine Similarity:

Cosine Similarity(A, B) = (A · B) / (||A|| × ||B||)

Score = 0.76 → The sentences are similar!
(Even though they share no common words)

Demo: Similarity Calculation with Databricks

A. Using the ai_query function in SQL:

# Databricks notebook source
embeddingsDF = spark.sql("""
    SELECT
        -- Embedding for the first sentence
        ai_query(
            "databricks-gte-large-en",            -- Model name
            "XBox is an excellent gaming console"  -- Text
        ) AS embedding1,

        -- Embedding for the second sentence
        ai_query(
            "databricks-gte-large-en",
            "I can spend all day with PlayStation"
        ) AS embedding2
""")

display(embeddingsDF)

B. Cosine Similarity calculation:

import numpy as np

def cosine_similarity(embedding1, embedding2):
    # Convert lists to numpy arrays
    arr1 = np.array(embedding1)
    arr2 = np.array(embedding2)

    # Calculate cosine similarity
    cosineSimilarityScore = np.dot(arr1, arr2) / (
        np.linalg.norm(arr1) * np.linalg.norm(arr2)
    )

    return cosineSimilarityScore

# Extract embeddings and calculate the score
embeddingsRow = embeddingsDF.first()
embedding1 = embeddingsRow['embedding1']
embedding2 = embeddingsRow['embedding2']

score = cosine_similarity(embedding1, embedding2)
print(f"Cosine similarity score: {score}")
# Result: ~0.73 (similar sentences!)

C. Call via Python SDK (WorkspaceClient):

from databricks.sdk import WorkspaceClient

client = WorkspaceClient()

# Embedding 1 — user question
response1 = client.serving_endpoints.query(
    name="databricks-gte-large-en",
    input="How many days are covered under pre-hospitalization?"
)
embedding1 = response1.data[0].embedding

# Embedding 2 — text extracted from the insurance policy
response2 = client.serving_endpoints.query(
    name="databricks-gte-large-en",
    input="We will not be liable to pay Pre-hospitalization Medical Expenses "
         "for more than 90 days immediately preceding the Insured Person's admission..."
)
embedding2 = response2.data[0].embedding

score = cosine_similarity(embedding1, embedding2)
print(f"Cosine similarity score: {score}")
# High score → The question is very similar to the policy text

D. Built-in SQL function ai_similarity:

SELECT ai_similarity(
    'I can spend all day with PlayStation',
    'XBox is an excellent gaming console'
)
-- Databricks automatically uses the most recent model

Content Types Supported by Embeddings

┌──────────────────────────────────────────────────────────────────────┐
│                  SUPPORTED EMBEDDING TYPES                           │
├──────────────┬───────────────────────────────────────────────────────┤
│ 📝 Text      │ Grounding data, classification, clustering            │
├──────────────┼───────────────────────────────────────────────────────┤
│ 💻 Code      │ Duplicate detection, code assistants, search          │
├──────────────┼───────────────────────────────────────────────────────┤
│ 🖼️ Image    │ Image search, facial recognition                      │
├──────────────┼───────────────────────────────────────────────────────┤
│ 🔊 Audio     │ Voice similarity, speaker identification              │
├──────────────┼───────────────────────────────────────────────────────┤
│ 🎬 Video     │ Movie recommendations, content similarity             │
├──────────────┼───────────────────────────────────────────────────────┤
│ 🔀 Multimodal│ Image search by text, visual Q&A                      │
└──────────────┴───────────────────────────────────────────────────────┘

Available Embedding Models

Criteria for Selecting an Embedding Model

graph TD
    A[Choose an embedding model] --> B{Domain complexity}
    A --> C{Required latency}
    A --> D{Multilingual needs}
    A --> E{Embedding dimensions}
    A --> F{Multimodal content?}
    A --> G{Cost}

    B --> B1[Generic vs specialized]
    C --> C1[Larger models = slower]
    D --> D1[Check supported languages]
    E --> E1[More dimensions = more precision]
    F --> F1[Look for multimodal models]
    G --> G1[Compare pricing]

3.2 Preparing, Chunking, and Storing Data in a Delta Table

Preparation Flow Overview

flowchart LR
    A["📁 Databricks Volume\nContoso Policy.pdf"] --> B[READ_FILES]
    B --> C[ai_parse_document]
    C --> D{Chunking\nStrategy}
    D --> E["Format-specific\nChunking"]
    D --> F["Recursive\nChunking"]
    E --> G["(Delta Table\nDocumentChunks)"]
    F --> G
    G --> H["(Delta Table\nProcessedFiles)"]

Step 1: Creating Delta Tables

-- Table to store document chunks
CREATE TABLE IF NOT EXISTS insurance.rag.DocumentChunks
(
    Id          BIGINT    GENERATED BY DEFAULT AS IDENTITY,
    Text        STRING,
    Source      STRING,
    PolicyName  STRING,
    UpdatedOn   TIMESTAMP
)
TBLPROPERTIES( delta.enableChangeDataFeed = true )
-- ⚠️ IMPORTANT: enableChangeDataFeed is required for incremental
--               synchronization with Vector Search

-- Table to track processed files
CREATE TABLE IF NOT EXISTS insurance.rag.ProcessedFiles
(
    FilePath   STRING,
    UpdatedOn  TIMESTAMP
)
TBLPROPERTIES( delta.enableChangeDataFeed = true )

Step 2: Installing Dependencies

# Install required libraries
%pip install langchain langchain-text-splitters

dbutils.library.restartPython()

Step 3: Identifying Unprocessed Files

-- Create a view to find files not yet processed
CREATE TEMPORARY VIEW vwFilesToProcess
AS
    SELECT allFiles.path,
           allFiles._metadata.file_name AS filename

    FROM READ_FILES('/Volumes/insurance/rag/raw/', format => 'binaryFile') AS allFiles

        LEFT ANTI JOIN insurance.rag.ProcessedFiles AS processedFiles
            ON allFiles.path = processedFiles.FilePath;

-- Check which files need processing
SELECT * FROM vwFilesToProcess;

Step 4: Extracting Content with ai_parse_document

-- Create a view with the parsed content of unprocessed files
CREATE TEMPORARY VIEW vwParseDocuments
AS
    SELECT ai_parse_document(content) AS parsed_document
         , filesToProcess.path
         , filesToProcess.filename

    FROM READ_FILES('/Volumes/insurance/rag/raw/', format => 'binaryFile') AS allFiles

        JOIN vwFilesToProcess AS filesToProcess ON allFiles.path = filesToProcess.path;

-- Check parsed content
-- ai_parse_document returns an array of elements:
-- { type: "page_header" | "section_header" | "text", content: "..." }
SELECT * FROM vwParseDocuments;

Step 5: Chunking Strategies

Available strategies:

┌─────────────────────────────────────────────────────┐
│              CHUNKING STRATEGIES                     │
├──────────────────────┬──────────────────────────────┤
│ Fixed-size           │ Fixed character size          │
│ Overlapping          │ With overlap                  │
│ Sentence/Paragraph   │ Based on sentences/paragraphs │
│ Format-specific      │ Markdown/HTML elements        │
│ Structure-aware      │ Structured sections           │
│ Semantic             │ Based on meaning              │
│ Recursive            │ Hierarchical separators       │
└──────────────────────┴──────────────────────────────┘

Option A: Format-specific chunking (direct use of ai_parse_document):

-- Extract each document element as a separate row
SELECT path
     , item.content AS content
     , item.bbox.page_id[0] AS pageId

FROM
(
    SELECT transform(
                parsed_document:document.elements
                    ::ARRAY<STRUCT<content:STRING, bbox:ARRAY<STRUCT<page_id:INT>>>>
               , x -> x
           ) AS content
         , path

    FROM vwParseDocuments
)

LATERAL VIEW explode(content) exploded_table AS item
LIMIT 10

Option B: Recursive chunking with LangChain (recommended):

from pyspark.sql.types import *
from pyspark.sql.functions import *
from langchain_text_splitters import RecursiveCharacterTextSplitter

# Configure the text splitter
textSplitter = RecursiveCharacterTextSplitter(
    chunk_size    = 1000,   # Max size of each chunk (characters)
    chunk_overlap = 200,    # Overlap to preserve context

    separators    = ["\n\n", "\n", " ", ""],  # Separator priority order
    length_function = len
)

# Splitting function
def textSplitRecursive(text):
    if text is None:
        return []
    return textSplitter.split_text(text)

# Register as UDF for use in SQL
spark.udf.register("textSplitRecursiveSql", textSplitRecursive, ArrayType(StringType()))

Step 6: Inserting Chunks into the Delta Table

-- Insert chunks into DocumentChunks
INSERT INTO insurance.rag.DocumentChunks (Text, Source, PolicyName, UpdatedOn)

SELECT  explode( textSplitRecursiveSql(content) ) AS Text
      , path AS Source
      , filename AS PolicyName
      , CURRENT_TIMESTAMP AS UpdatedOn

FROM
(
    -- Combine all elements of a document into a single text
    SELECT array_join(
               transform(
                   parsed_document:document.elements
                       ::ARRAY<STRUCT<content:STRING>>, x -> x.content
               ),
               '\n'
           ) AS content
         , path
         , filename

    FROM vwParseDocuments
)

-- Result: 264 chunks created for a 52-page document

Step 7: Marking Files as Processed

-- Insert processed files to avoid reprocessing
INSERT INTO insurance.rag.ProcessedFiles (FilePath, UpdatedOn)
SELECT path, CURRENT_TIMESTAMP FROM vwFilesToProcess;

-- Verify
SELECT * FROM insurance.rag.ProcessedFiles

💡 Why is overlap important?
A chunk_overlap of 200 characters ensures context is not lost due to chunking. Adjacent chunks share a common excerpt, which improves search accuracy.

3.3 Setting Up Mosaic AI Vector Search

Overview of Databricks Vector Search

┌─────────────────────────────────────────────────────────────────────┐
│               DATABRICKS VECTOR SEARCH (Mosaic AI)                  │
├─────────────────────────────────────────────────────────────────────┤
│                                                                      │
│  ✅ Serverless architecture (no infrastructure provisioning)         │
│  ✅ Native integration with Delta Tables (Unity Catalog)             │
│  ✅ Automatic Delta Table → Index synchronization                    │
│  ✅ Governance via Unity Catalog (security, lineage)                 │
│  ✅ High scalability (billions of vectors, low latency)              │
│  ✅ UI interface, Python SDK, REST API                               │
│                                                                      │
│  CAPABILITIES:                                                       │
│  • Standard  : 320 million vectors (768 dim)                        │
│  • Storage   : >1 billion vectors (cheaper, higher latency)          │
└─────────────────────────────────────────────────────────────────────┘

Supported Search Types

graph TD
    VS[Vector Search] --> KW["Keyword Search\nBM25 — exact words"]
    VS --> SIM["Similarity Search\nCosine Similarity + ANN"]
    VS --> HYB["Hybrid Search\nKeyword + Similarity\n→ Reciprocal Rank Fusion"]
    VS --> META["Metadata Filtering\nSQL filters on columns"]

    HYB --> RRF["RRF — Reciprocal Rank Fusion\nCombines results"]

Setting Up the Vector Search Endpoint

Via the Databricks interface:

1. Left menu → Compute → Vector Search
2. Create a new endpoint: insurancevectorsearch
3. Choose type: Standard or Storage Optimized
4. Create

⚠️ Cost warning: As soon as an endpoint is created, billing begins. Delete the endpoint if not in use.

3.4 Indexing and Querying Data in Vector Search

Vector Search Index Structure

┌─────────────────────────────────────────────────────────────────────┐
│                    VECTOR SEARCH INDEX                               │
│                                                                      │
│  Specialized structure that organizes embedding vectors             │
│  to enable fast similarity search WITHOUT scanning                  │
│  every vector.                                                       │
│                                                                      │
│  Stored columns:                                                     │
│  ┌──────────┬──────────┬────────────┬──────────┬───────────────┐    │
│  │    Id    │   Text   │   Source   │PolicyName│  Embedding[]  │    │
│  ├──────────┼──────────┼────────────┼──────────┼───────────────┤    │
│  │   001    │ "90 days │ /path/...  │ Contoso  │[0.17, -1.75,] │    │
│  │   002    │ "Covered │ /path/...  │ Contoso  │[3.29, 1.18, ] │    │
│  │   ...    │   ...    │   ...      │  ...     │    ...        │    │
│  └──────────┴──────────┴────────────┴──────────┴───────────────┘    │
└─────────────────────────────────────────────────────────────────────┘

Delta Table → Index Sync Options

┌──────────────────────────────────────────────────────────────────────┐
│              SYNCHRONIZATION OPTIONS                                 │
├───────────────────────────────┬──────────────────────────────────────┤
│ Delta Sync Index              │ Delta table: data only               │
│ (auto-generated embeddings)   │ Databricks generates embeddings      │
│                               │ during automatic sync                │
├───────────────────────────────┼──────────────────────────────────────┤
│ Delta Sync Index              │ Delta table: data + pre-computed     │
│ (self-managed embeddings)     │ embeddings                           │
│                               │ Databricks auto-syncs data           │
├───────────────────────────────┼──────────────────────────────────────┤
│ Direct Vector Access Index    │ Manual embedding management          │
│ (manual)                      │ AND synchronization via REST API     │
└───────────────────────────────┴──────────────────────────────────────┘

Creating the Index from the Interface

1. Catalog → documentchunks Table → Create → Vector Search Index
2. Properties:
   • Name: insuranceindex
   • Primary key: Id column
   • Columns: all (leave empty)
   • Embedding generation: on Text column
   • Embedding model: databricks-gte-large-en
   • Endpoint: insurancevectorsearch
3. Create → Monitor the synchronization pipeline

Complete Data Preparation Flow

flowchart LR
    VOL["Volume\nContoso.pdf"] --> READ[READ_FILES]
    READ --> PARSE[ai_parse_document]
    PARSE --> SPLIT["RecursiveCharacterTextSplitter\nchunk_size=1000, overlap=200"]
    SPLIT --> DT["(Delta Table\nDocumentChunks)"]
    DT -->|Delta Sync| EMB["Embedding Model\ndatabricks-gte-large-en"]
    EMB --> IDX["(Vector Search Index\ninsuranceindex)"]

    style VOL fill:#fff3e0
    style DT fill:#e3f2fd
    style IDX fill:#e8f5e9

4. Module 3 — Building the RAG Pipeline with Databricks

4.1 Understanding the RAG Workflow in Databricks

Complete Workflow Overview

flowchart TD
    subgraph PREP["Data Preparation (once)"]
        D[Documents] --> TS[Text Splitter] --> C[Chunks]
        C --> EM1[Embedding Model] --> VI[(Vector Search Index)]
    end

    subgraph QUERY["Query Response (real-time)"]
        U[👤 User] -->|Question| EM2[Embedding Model]
        EM2 -->|Question Vector| VS["Vector Search\nCosine Similarity"]
        VS -->|Top-K Chunks| AUG["Augmentation\nQuestion + Chunks + Instructions"]
        AUG --> LLM["LLM\ndatabricks-gpt-oss-120b"]
        LLM -->|Grounded response| U
    end

    VI --> VS

    style PREP fill:#e3f2fd
    style QUERY fill:#e8f5e9

Steps for a Production RAG

┌─────────────────────────────────────────────────────────────────────┐
│              STEPS FOR A PRODUCTION RAG                              │
├─────────────────────────────────────────────────────────────────────┤
│                                                                      │
│  1. BUILD the RAG pipeline                                           │
│     ├─ Similarity search (Vector Search)                            │
│     ├─ Prompt preparation (augmentation)                             │
│     └─ Response generation (LLM)                                    │
│                                                                      │
│  2. TRACK the pipeline with MLflow                                   │
│     ├─ Log parameters, metrics                                       │
│     └─ Trace each step (chunks, prompt, response)                   │
│                                                                      │
│  3. EVALUATE the pipeline                                            │
│     ├─ Test prompts                                                  │
│     └─ Scores: correctness, relevance, safety                       │
│                                                                      │
│  4. REGISTER the RAG model in Unity Catalog (MLflow)                │
│     └─ Model versioning                                              │
│                                                                      │
│  5. DEPLOY the model to an endpoint                                  │
│     └─ URL accessible by any application                            │
│                                                                      │
└─────────────────────────────────────────────────────────────────────┘

4.2 Prompt Engineering for RAG

Structure of an Effective RAG Prompt

┌─────────────────────────────────────────────────────────────────────┐
│                    RAG PROMPT TEMPLATE                               │
├─────────────────────────────────────────────────────────────────────┤
│                                                                      │
│  [SYSTEM PROMPT]                                                     │
│  You are a health insurance assistant answering questions            │
│  using the provided context.                                         │
│                                                                      │
│  Context:                                                            │
│  {context_chunks}          ← Chunks retrieved from Vector Search     │
│                                                                      │
│  Question:                                                           │
│  {user_question}           ← User's question                        │
│                                                                      │
│  Instructions:                                                       │
│  - Answer using only the context above                              │
│  - If answer is not in context, say "Information not available"     │
│  - Provide citations for sources in [source] format                 │
│                                                                      │
└─────────────────────────────────────────────────────────────────────┘

Context Injection Patterns

┌─────────────────┬──────────────────────┬────────────────────────────┐
│ SIMPLE          │ STRUCTURED           │ DELIMITED                  │
├─────────────────┼──────────────────────┼────────────────────────────┤
│ Context:        │ Context:             │ <context>                  │
│ {chunk1}        │ [Source: <path>      │ Context:                   │
│ {chunk2}        │  | Page: <page>]     │ {chunk1}                   │
│ {chunk3}        │ {chunk1}             │ {chunk2}                   │
│                 │ [Source: <path>      │ {chunk3}                   │
│ Simple but      │  | Page: <page>]     │ </context>                 │
│ may confuse     │ {chunk2}             │                            │
│ the model       │ Citations included   │ Clear boundaries           │
└─────────────────┴──────────────────────┴────────────────────────────┘

Essential Prompt Instructions

4.3 Building the RAG Pipeline

Installing Dependencies

%pip install databricks-vectorsearch mlflow
%restart_python

Defining Variables and Clients

import mlflow
from databricks.vector_search.client import VectorSearchClient
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

# Vector Search endpoint & index
vectorSearchEndpointName = "insurancevectorsearch"
indexName                = "insurance.rag.insuranceindex"

# LLM model
llmModelName = "databricks-gpt-oss-120b"

# Databricks workspace URL
workspaceUrl = "https://" + spark.conf.get("spark.databricks.workspaceUrl")

# Personal access token
# ⚠️ In production: use Databricks Secrets
accessToken = "***ACCESS TOKEN***"

# Create the Vector Search client
client = VectorSearchClient(
    workspace_url         = workspaceUrl,
    personal_access_token = accessToken,
    disable_notice        = True
)

# Get a reference to the index
index = client.get_index(
    endpoint_name = vectorSearchEndpointName,
    index_name    = indexName
)

# Workspace client to call the LLM
workspace = WorkspaceClient(
    host  = workspaceUrl,
    token = accessToken
)

💡 Note: In a Databricks notebook, workspace_url and personal_access_token can be omitted (auto-detected). They are required when deploying the model.

Step 1: RETRIEVAL — Fetching Chunks

def retrieveChunks(query):
    results = index.similarity_search(
        query_text  = query,                              # Query text

        query_type  = "HYBRID",                           # Options:
                                                          # FULL_TEXT: keyword search
                                                          # ANN      : semantic search (default)
                                                          # HYBRID   : keyword + semantic

        columns     = ["Text", "Source", "PolicyName"],   # Columns to retrieve

        num_results = 3,                                  # Top-k strategy

        score_threshold = 0.8,                            # Minimum score

        filters     = {"PolicyName LIKE": "Contoso"}      # SQL filters
    )
    return results

Step 2: Context Preparation

def prepareContext(chunks):
    docs = ([
        "[Source: " + r[1] + "] \n[Content]: " + r[0]
            for r in chunks["result"]["data_array"]
    ])
    chunksContext = "\n\n".join(docs)
    return chunksContext

Step 3: AUGMENTATION — Building the Prompt

# System message
systemMessage = "You are a health insurance assistant answering questions using the provided context."

# Prompt template
PROMPT_TEMPLATE = """
Context:
{context}

Question:
{question}

Instructions:
- Answer using only the context above
- If answer is not contained in the context, say "Information not available"
- Provide citations for sources in [Source - Content] format
"""

def buildPrompt(query, chunksContext):
    prompt = PROMPT_TEMPLATE.format(
        context  = chunksContext,
        question = query
    )
    return prompt

Step 4: GENERATION — LLM Call

def generateAnswer(systemMessage, prompt):
    # Call the pre-deployed model
    response = workspace.serving_endpoints.query(
        name = llmModelName,   # databricks-gpt-oss-120b
        messages = [
            ChatMessage(
                role    = ChatMessageRole.SYSTEM,
                content = systemMessage,
            ),
            ChatMessage(
                role    = ChatMessageRole.USER,
                content = prompt,
            )
        ]
    )
    # Extract the response
    answer = response.choices[0].message.content[1]["text"]
    return answer

Complete RAG Pipeline

def ragPipeline(query):
    # 1. RETRIEVAL
    chunks = retrieveChunks(query)

    # 2. Context preparation
    chunksContext = prepareContext(chunks)

    # 3. AUGMENTATION
    prompt = buildPrompt(query, chunksContext)

    # 4. GENERATION
    answer = generateAnswer(systemMessage, prompt)

    return {"predictions": [{"content": answer}]}

# Test
input = {"messages": [{"role": "user", "content": "How many days are covered under pre-hospitalization?"}]}
query = input["messages"][0]["content"]

chunks = retrieveChunks(query)
chunksContext = prepareContext(chunks)
prompt = buildPrompt(query, chunksContext)
answer = generateAnswer(systemMessage, prompt)

# Result: "Pre-hospitalization is covered up to 90 days [Source: ...]"

RAG Processing Flow

User question
        │
        ▼
┌───────────────────┐
│ retrieveChunks()  │  → similarity_search() on Vector Search
│ (RETRIEVAL)       │  → Returns top-3 chunks with score ≥ 0.8
└───────────────────┘
        │
        ▼
┌───────────────────┐
│ prepareContext()  │  → Formats chunks: [Source: ...] [Content: ...]
│ (CONTEXT)         │
└───────────────────┘
        │
        ▼
┌───────────────────┐
│ buildPrompt()     │  → Injects context + question + instructions
│ (AUGMENTATION)    │  into PROMPT_TEMPLATE
└───────────────────┘
        │
        ▼
┌───────────────────┐
│ generateAnswer()  │  → LLM call (GPT-OSS-120B) via serving_endpoints
│ (GENERATION)      │  → Returns response with citations
└───────────────────┘
        │
        ▼
Response grounded in enterprise data

4.4 Tracking and Evaluating the RAG Pipeline with MLflow

MLflow Overview

MLflow is an open-source platform that manages the end-to-end lifecycle of machine learning and GenAI applications.

┌─────────────────────────────────────────────────────────────────────┐
│                    MLFLOW FOR RAG                                    │
├───────────────────────────┬─────────────────────────────────────────┤
│ 📊 Track experiments      │ Log prompts, parameters, metrics        │
├───────────────────────────┼─────────────────────────────────────────┤
│ 🔬 Evaluate quality       │ Scorers: correctness, relevance, safety │
├───────────────────────────┼─────────────────────────────────────────┤
│ 🔍 Trace pipelines        │ Inspect output at each step             │
├───────────────────────────┼─────────────────────────────────────────┤
│ 📦 Manage & deploy models │ Register, version, deploy               │
└───────────────────────────┴─────────────────────────────────────────┘

Installing Dependencies

%pip install databricks-agents
%pip install databricks-sdk --upgrade
%pip install mlflow[databricks]

dbutils.library.restartPython()

Creating an MLflow Experiment

import mlflow

# Define the experiment name
experimentName = "/Workspace/Databricks RAG/RAG Experiment"

# Create the experiment (once)
experimentId = mlflow.create_experiment(experimentName)

Tracking and Tracing

from mlflow.models import infer_signature

# Define parameters
numOfResults   = 3
scoreThreshold = 0.8

# Activate the experiment
mlflow.set_experiment(experimentName)

# Start an MLflow run
with mlflow.start_run(run_name="rag-pipeline-run") as run:

    input = {"messages": [{"role": "user", "content": "What is the coverage under permanent loss of speech?"}]}
    query = input["messages"][0]["content"]

    # Execute the RAG pipeline
    chunks        = retrieveChunks(query)
    chunksContext = prepareContext(chunks)
    prompt        = buildPrompt(query, chunksContext)
    answer        = generateAnswer(systemMessage, prompt)

    print(answer)

    # ── TRACKING: log parameters ──────────────────────────────────
    mlflow.log_param("model",           llmModelName)
    mlflow.log_param("top_k",           numOfResults)
    mlflow.log_param("score_threshold", scoreThreshold)

    # ── TRACKING: log metrics ──────────────────────────────────────
    metrics = {"retrieved_docs": len(chunks)}
    mlflow.log_metrics(metrics)

    # ── TRACING: log output of each step ──────────────────────────
    mlflow.log_text(query,         "query.txt")
    mlflow.log_text(chunksContext, "context.txt")
    mlflow.log_text(prompt,        "prompt.txt")
    mlflow.log_text(answer,        "response.txt")
    mlflow.log_dict({"citations": chunks}, "citations.json")

    runId = run.info.run_id
    print(f"Run Id: {runId}")

Evaluating the RAG Pipeline

import mlflow
from mlflow.genai.scorers import RelevanceToQuery, Safety, Correctness
import pandas as pd

mlflow.set_experiment(experimentName)

# Prediction method for evaluation
def predict(messages):
    query = messages[0]["content"]
    answer = ragPipeline(query)
    return answer

# Evaluation data with expected results
data = [
    {
        "inputs": {
            "messages": [{"role": "user", "content": "What is the coverage under permanent loss of speech?"}]
        },
        "expectations": {"expected_facts": ["70%"]}
    },
    {
        "inputs": {
            "messages": [{"role": "user", "content": "How many days are covered under pre-hospitalization?"}]
        },
        "expectations": {"expected_facts": ["90 days"]}
    },
    {
        "inputs": {
            "messages": [{"role": "user", "content": "What policies are available from Contoso?"}]
        },
        "expectations": {"expected_facts": ["contoso premier policy"]}
    }
]

# Launch evaluation
results = mlflow.genai.evaluate(
    data       = data,
    predict_fn = predict,
    scorers    = [
        RelevanceToQuery(),  # Check if the response is relevant to the query
        Safety(),            # Check for harmful content
        Correctness()        # Check match with expected facts
    ]
)

Available MLflow Scorers:

4.5 Registering and Deploying the RAG Model

Defining the Python Model Class

# Class inheriting from mlflow.pyfunc.PythonModel
class InsuranceRagModel(mlflow.pyfunc.PythonModel):

    def predict(self, context, model_input):
        # Input is received as a Pandas DataFrame during deployment
        if isinstance(model_input, pd.DataFrame):
            query = model_input["messages"].iloc[0][0]["content"]
        else:
            query = model_input["messages"][0]["content"]

        # Call the RAG pipeline
        answer = ragPipeline(query)

        return answer

# Set as active/current model
mlflow.models.set_model(model = InsuranceRagModel())

Registering the Model in Unity Catalog

from mlflow.models import infer_signature

mlflow.set_experiment(experimentName)

with mlflow.start_run(run_name="rag-pipeline-model-register") as run:

    modelInfo = mlflow.pyfunc.log_model(
        name = "InsuranceRagModel",                         # Model name in registry

        python_model = "RAG Model",                         # Notebook containing RAG components

        registered_model_name = "insurance.rag.insurancemodel",  # Path in Unity Catalog

        signature = infer_signature(input, answer)          # Input/output format
    )

print(f"MLflow Run: {modelInfo.run_id}")
print(f"Model URI: {modelInfo.model_uri}")

Deploying to a Serving Endpoint

flowchart LR
    A["RAG Code\nNotebook"] -->|mlflow.pyfunc.log_model| B["(Unity Catalog\ninsurance.rag.insurancemodel)"]
    B -->|Create Serving| C["Serving Endpoint\ninsuranceModelEndpoint"]
    C -->|REST API| D[Client Application]

    style B fill:#e3f2fd
    style C fill:#e8f5e9

Endpoint configuration:

• Name: insuranceModelEndpoint
• Model: insurance.rag.insurancemodel (version automatically selected)
• Compute type: Choose machine size
• Scale-out: Configure auto-scaling (can scale back to 0 machines)
• Tracing: Enable to capture logs

Testing the endpoint:

// Required input format
{
  "messages": [
    {
      "role": "user",
      "content": "How many days are covered under pre-hospitalization?"
    }
  ]
}

5. Module 4 — Agent Bricks for RAG Workflows

5.1 Understanding Agent Bricks in Databricks Mosaic AI

Direct LLM vs Agent

┌────────────────────────────────────────────────────────────────┐
│  DIRECT LLM (classic RAG)                                      │
│  Input → [GenAI Model] → Output                                │
│  • Single LLM call                                             │
│  • Simple response                                             │
├────────────────────────────────────────────────────────────────┤
│  AGENT                                                         │
│  Input → [GenAI Model + Tools + Instructions] → Output         │
│  • Autonomous reasoning                                        │
│  • Multiple tool calls                                         │
│  • Multi-step resolution                                       │
└────────────────────────────────────────────────────────────────┘

Agent Architecture

graph TD
    U[👤 User] -->|Complex question| A[AGENT]

    subgraph AGENT["Agent (Mosaic AI)"]
        LLM["LLM\ndatabricks-gpt-oss-120b"]
        TOOLS[Tools]
        INST["Instructions\nPersona"]
    end

    TOOLS --> VS["Vector Search\nPolicy details"]
    TOOLS --> SQL["SQL Function\nClaim amount"]
    TOOLS --> CODE["Python Code\nExecutor"]

    A --> R["Reasoning\n& Planning"]
    R --> STEP1["Step 1: Check\ncovered amount"]
    R --> STEP2["Step 2: Check\nexclusions"]
    STEP1 --> |Tool call| SQL
    STEP2 --> |Tool call| VS
    STEP1 --> COMBINE[Combine results]
    STEP2 --> COMBINE
    COMBINE --> U

    style AGENT fill:#fff3e0

Agent Capabilities

┌─────────────────────────────────────────────────────────────────────┐
│                    AGENT CAPABILITIES                               │
├─────────────────────────────────────────────────────────────────────┤
│                                                                      │
│  🧠 REASONING        Understand the problem and decide              │
│                       next steps                                    │
│                                                                      │
│  📋 PLANNING         Break tasks into smaller steps                 │
│                                                                      │
│  🔧 TOOL USE         Call functions, query databases,               │
│                       perform vector searches                       │
│                                                                      │
│  💾 MEMORY           Maintain previous conversations                │
│                                                                      │
│  🔄 AUTONOMOUS       Multi-step workflows and iteration             │
│      EXECUTION       until the goal is achieved                     │
│                                                                      │
└─────────────────────────────────────────────────────────────────────┘

Example: Insurance Agent

User question:
"How much is my coverage for cosmetic surgery, and what is excluded?"

─────────────────────────────────────────────────────────────────────
AGENT REASONING:
  → Two distinct questions to resolve
  → Step 1: Coverage amount for cosmetic surgery
  → Step 2: Exclusions related to cosmetic surgery

TOOLS CALLED:
  → [SQL Function]: Retrieves covered amount from claims table
  → [Vector Search]: Searches exclusions in the insurance policy

FINAL RESPONSE:
  Combination of both results into a coherent response
─────────────────────────────────────────────────────────────────────

Agentic RAG vs Classic RAG

Agent Bricks in Mosaic AI

graph TD
    AB[Agent Bricks] --> RE[Reasoning Engine]
    AB --> MEM[Memory]
    AB --> TC[Tool Calling]
    AB --> RET[Retriever]
    AB --> ORCH[Orchestrator]

    AB --> INT[Integrations]
    INT --> UC["Unity Catalog\nData & Security"]
    INT --> VS["Vector Search\nSemantic retrieval"]
    INT --> MLF["MLflow\nTracking & Evaluation"]
    INT --> COMP["Managed Compute\nDeployment"]
    INT --> GW["AI Gateway\nGuardrails"]

    AB --> APPR[Approaches]
    APPR --> LC["Low-code\nTemplates"]
    APPR --> CF["Code-first\nDatabricks Agents SDK"]

Pre-built templates in Agent Bricks:

5.2 Building an Agentic RAG with Agent Bricks

Three Options for Building an Agentic RAG

┌──────────────────────┬─────────────────────────────┬──────────────────────────────┐
│ OPTION 1             │ OPTION 2                    │ OPTION 3                     │
│ Playground UI        │ Knowledge Assistant         │ Databricks Agents SDK        │
│ + Vector Search      │ Template                    │                              │
├──────────────────────┼─────────────────────────────┼──────────────────────────────┤
│ • Rapid prototyping  │ • Faster development        │ • Full control               │
│ • Code generation    │ • Optimized architecture    │ • Complete customization     │
│ • Interactive testing│   (Instructor-Receiver)     │ • Agentic RAG from scratch   │
│                      │ • No code required          │                              │
└──────────────────────┴─────────────────────────────┴──────────────────────────────┘

Option 1: Playground UI

Steps:

1. Left menu → Playground
2. Select the model endpoint (e.g., pre-deployed GPT model)
3. Add a tool → Vector Search Index → select insuranceindex
4. Optional: Define the system prompt and instructions
5. Test the question: “How much is the coverage for cosmetic surgery and what is excluded?”

Agentic RAG observation:

• The agent calls the index multiple times (once for coverage, once for exclusions)
• Visible in the Trace: multiple calls to the Vector Search tool
• Possible to generate code from the Playground to reproduce the agent

Option 2: Knowledge Assistant Template

Instructor-Receiver architecture (different from classic RAG):

flowchart LR
    U[👤 User] --> IA["Instructor Agent\n(orchestrating LLM)"]
    IA --> RA["Receiver Agent\n(specialized RAG)"]
    RA --> VS["(Vector Search\nIndex)"]
    VS -->|Chunks| RA
    RA -->|Structured response| IA
    IA -->|Final response| U

    style IA fill:#e3f2fd
    style RA fill:#e8f5e9

This architecture is more optimized than traditional RAG for Q&A.

Configuration steps:

1. Left menu → Agents → Create a new agent
2. Select the Knowledge Assistant template
3. Configure:
   • Name: InsuranceKnowledgeAssistant
   • Description: Insurance agent
   • System message: Instructions for the model
4. Add a knowledge source:
   • Type: Vector Search Index
   • Source: insuranceindex
   • Doc URI: Source column (source file path)
   • Chunk column: Text column
   • Description: “This source contains insurance policy details” (crucial for the agent to know when to use this tool)
5. Optional: custom instructions
6. Create the agent

After creation, you get:

┌─────────────────────────────────────────────────────────────────────┐
│  INSURANCE KNOWLEDGE ASSISTANT — Deployed                           │
├─────────────────────────────────────────────────────────────────────┤
│                                                                      │
│  🔐 Permissions   : Define who can query the agent                  │
│  🌐 Endpoint      : Managed Databricks endpoint (billed)            │
│  📊 MLflow Exp.   : Experiment created for tracking/tracing          │
│  🧪 Test in UI    : Built-in testing interface                      │
│                                                                      │
│  Test: "How much is the coverage for cosmetic surgery?"              │
│  → Response with citations and visible trace                        │
│                                                                      │
└─────────────────────────────────────────────────────────────────────┘

Option 3: Databricks Agents SDK (full control)

For more control, use the SDK to:

• Build an Agentic RAG from scratch
• Use the functions from previous modules
• Create a Knowledge Assistant programmatically
• Fully customize agent behavior

6. Environment Setup

Prerequisites

Option 1: Azure Databricks

1. Sign up for Azure: https://azure.microsoft.com/en-in/free/
2. Go to the Azure portal: portal.azure.com
3. Search Azure Databricks → Create
4. Configuration:
   • Resource group: e.g., PluralsightDemos
   • Workspace name: unique
   • Region: e.g., East US 2
   • Pricing tier: Premium (required for Unity Catalog)
5. Launch the workspace

Option 2: Databricks Free Edition

2. Validate your email address
3. Configure your account

Unity Catalog Structure for This Course

insurance (catalog)
└── rag (schema)
    ├── DocumentChunks (Delta Table)
    │   ├── Id (BIGINT, auto-generated)
    │   ├── Text (STRING)
    │   ├── Source (STRING)
    │   ├── PolicyName (STRING)
    │   └── UpdatedOn (TIMESTAMP)
    ├── ProcessedFiles (Delta Table)
    │   ├── FilePath (STRING)
    │   └── UpdatedOn (TIMESTAMP)
    ├── insuranceindex (Vector Search Index)
    └── insurancemodel (MLflow Model)

Volumes:
└── insurance.rag.raw
    └── Contoso Premier Policy.pdf

7. Code File Reference

Python Dependencies

# Module 2
%pip install langchain langchain-text-splitters

# Module 3
%pip install databricks-vectorsearch mlflow

# Module 4
%pip install databricks-agents
%pip install databricks-sdk --upgrade
%pip install mlflow[databricks]

Models Used

Summary — Complete RAG Architecture with Databricks

flowchart TD
    subgraph PREP["📦 Data Preparation (once)"]
        A["Contoso Premier\nPolicy.pdf"] --> B["Databricks Volume\ninsurance.rag.raw"]
        B --> C["READ_FILES\nai_parse_document"]
        C --> D["RecursiveCharacterTextSplitter\nchunk_size=1000, overlap=200"]
        D --> E["(Delta Table\nDocumentChunks)"]
        E --> F["Delta Sync\ndatabricks-gte-large-en"]
        F --> G["(Vector Search Index\ninsuranceindex)"]
    end

    subgraph PIPELINE["🔄 RAG Pipeline (real-time)"]
        H[👤 User] -->|Question| I["retrieveChunks\nHYBRID Search"]
        I --> G
        G -->|Top-3 chunks| J[prepareContext]
        J --> K["buildPrompt\nContext + Question + Instructions"]
        K --> L["generateAnswer\ndatabricks-gpt-oss-120b"]
        L -->|Response with citations| H
    end

    subgraph MLOPS["⚙️ MLOps with MLflow"]
        L --> M["MLflow Tracking\nParameters + Metrics"]
        L --> N["MLflow Evaluation\nRelevance + Safety + Correctness"]
        L --> O["(Unity Catalog\ninsurance.rag.insurancemodel)"]
        O --> P["Serving Endpoint\ninsuranceModelEndpoint"]
    end

    subgraph AGENT["🤖 Agentic RAG (Agent Bricks)"]
        Q[👤 User] -->|Complex question| R["Agent\nInsuranceKnowledgeAssistant"]
        R --> G
        R --> S["SQL Function\nClaim amounts"]
        R -->|Multi-step response\nwith reasoning| Q
    end

    style PREP fill:#e3f2fd,stroke:#1565c0
    style PIPELINE fill:#e8f5e9,stroke:#2e7d32
    style MLOPS fill:#fff3e0,stroke:#e65100
    style AGENT fill:#fce4ec,stroke:#880e4f

Search Terms

rag · pipelines · databricks · vector · search · embeddings · artificial · intelligence · generative · ai · agent · option · pipeline · bricks · delta · dependencies · model · mosaic · prompt · agentic · data · environment · flow · index