RAG on Databricks

Fabric Harness wires Databricks-native products for retrieval, generation, tracing, and evaluation. It does not reimplement AI Search, Unity AI Gateway, the offline index pipeline, or MLflow judges.

Follow the same mental model as the Databricks AI Cookbook RAG inference chain:

1. (Optional) preprocess the user query
2. Retrieve with Databricks AI Search
3. Augment the prompt with retrieved context
4. Generate through Unity AI Gateway or a custom Model Serving endpoint
5. Validate inline citations and apply answer limits

Offline chunk → embed → index remains a Databricks Job / notebook / Lakeflow concern. Quality measurement uses MLflow 3 evaluation and managed Databricks judges, not a parallel Fabric-only judge product.

Quick start (managed recipe)

fh add databricks rag-chain
# set DATABRICKS_HOST, OAuth credentials, DATABRICKS_AI_SEARCH_INDEX, DATABRICKS_MODEL
fh run rag-answer --question "How do I reset my password?"

For a credential-free local run, add --mock:

fh run rag-answer --question "How do I reset my password?" --mock

The generated recipe then selects its deterministic retriever and model fixture. This exercises discovery, input/output validation, the RAG chain, and citation handling without calling AI Search or Model Serving. Remove --mock to use the configured Databricks index and model.

Scaffolded files:

Related recipes:

Code API

Deterministic chain (cookbook online path)

import { databricksRagChain, toMlflow3EvaluationRecord, exportMlflow3EvaluationJsonl } from '@fabric-harness/databricks';

const chain = databricksRagChain({
  databricks: {
    host: process.env.DATABRICKS_HOST!,
    principal: { kind: 'pat', token: process.env.DATABRICKS_TOKEN! },
    model: 'system.ai.gpt-oss-20b',
    aiSearch: {
      index: 'main.support.kb_index',
      textColumn: 'chunk',
      idColumn: 'id',
      inputMode: 'text',
      strategy: 'hybrid',
    },
  },
  retrieval: { k: 5, scoreThreshold: 0.25 },
  postProcess: {
    requireCitations: true,
    citationRepairAttempts: 2,
  },
});

const turn = await chain.invoke({ question: 'How do I reset my password?' });
// turn.answer, turn.sources, turn.citations, turn.retrieval, turn.usage

With requireCitations, Fabric asks the model to revise an uncited answer or an answer that cites an unknown ID using only the exact source IDs returned by retrieval. Repair is bounded and fails closed if the model still omits a valid marker or retains a fabricated ID; Fabric never adds a citation to the answer on the model's behalf.

Complete AI Search query controls

Text input uses native HYBRID retrieval by default. Query representation and retrieval strategy are separate:

const result = await bundle.retriever.query('quarterly retention', {
  inputMode: 'text',
  strategy: 'hybrid',       // ann | hybrid | full-text
  k: 20,
  filter: { language: 'en' },
  scoreThreshold: 0.35,
  queryColumns: ['chunk', 'title'],
  sortColumns: ['published_at DESC'],
  facets: ['language', 'product'],
  reranker: { model: 'databricks_reranker' },
  columnsToRerank: ['title', 'chunk'],
  signal: abortController.signal,
});

result.chunks;
result.facetResult;
result.nextPageToken;
result.response; // complete generated-SDK response

Use inputMode: 'vector' with strategy: 'ann' for a self-managed embedding index. Supply an embeddingEndpoint when configuring the bundle or provide a precomputed queryVector for a specific request. text-and-vector sends both inputs for a supported native strategy.

The same controls can be defaults on databricksRagChain({ retrieval: ... }) or per-turn:

const turn = await chain.invoke({
  question: 'What changed?',
  retrieval: {
    strategy: 'hybrid',
    k: 12,
    facets: ['release'],
    reranker: { model: 'databricks_reranker' },
  },
});

Pagination and facet metadata are retained on RagTurn.retrieval. For direct search calls, retriever.nextPage(token, { signal }) retrieves the following page.

Streaming

chain.stream() yields incremental delta events while the model generates, then one terminal turn event carrying the same validated RagTurn that invoke returns:

for await (const event of chain.stream({ question: 'How do I reset my password?' })) {
  if (event.type === 'delta') process.stdout.write(event.textDelta);
  if (event.type === 'turn') persist(event.turn); // validated answer, citations, usage
}

Deltas are raw first-pass model output for live UI — they are emitted before citation validation, bounded repair, truncation, and the sources footer run, so the concatenated deltas can differ from the final answer. The terminal turn is authoritative: use it for persistence, evaluation export, and cost telemetry. A failed citation validation raises from the iterator after the deltas were observed. When the model provider does not implement stream(), the chain falls back to generate() and emits the whole answer as one delta.

The protected rag certification check uses this streaming path against the configured live AI Search index and Databricks model endpoint. It requires text deltas, a terminal turn, retrieved context, and at least one valid inline citation.

Release run 29622641870 certified the exact 2.0.0 package candidate in the protected Azure eastus2 workspace: the streaming chain emitted 12 text deltas, retrieved three source chunks, returned a valid fabric-databricks citation, and produced its authoritative terminal turn. The same run completed Databricks Job 613907318899842 for the managed 50-case RAG evaluation with a SUCCESS result. Agent Mode was not configured in that workspace and is not implied by this RAG evidence.

That record predates the 3.0 native-SDK transport and is historical for 3.0 promotion. The 3.0 release gate reruns the same retrieval, streaming, citation, and managed-evaluation checks against the exact package artifact.

What runs under the hood:

Agentic multi-tool RAG

When the agent must also call SQL, Genie, or Jobs, keep tool-calling:

const chain = databricksRagChain({ databricks: { /* + aiSearch */ } });
const { modelProvider, tools, policy } = chain.asAgentTools();
// pass into init({ modelProvider, tools, policy })

Or continue using databricks({ aiSearch }) + bundle.tools as in examples/with-databricks-rag.

Evaluation and quality (native Databricks first)

Local CI smoke (optional)

Lightweight checks on a RagTurn validate actual [source-id] markers, required facts, and retrieval:

import { scoreRagTurn, toMlflow3EvaluationRecord, exportMlflow3EvaluationJsonl } from '@fabric-harness/databricks';

const scores = scoreRagTurn(turn, {
  mustContain: ['Settings'],
  requireCitations: true,
  mustRetrieveIds: ['doc-1'],
});

These are smoke scorers, not a substitute for managed evaluation.

What developers inspect

Keep the answer and its source references visible beside quality evidence. The representative state below shows the relationship between retrieved context, citations, managed judges, latency, and cost; replace the fixture corpus and expected facts with the application's governed domain data.

MLflow 3 managed evaluation

Export MLflow 3 rows with structured inputs, outputs, expectations, retrieved context, and trace metadata:

const record = toMlflow3EvaluationRecord(turn, {
  expectedAnswer: 'Use Settings, then Security.',
  traceId: 'tr-...',
  submissionId: 'sub-...',
});
const jsonl = exportMlflow3EvaluationJsonl([record]);
// Merge the record into an MLflow Evaluation Dataset from a Databricks notebook or job.

The release certification fixture runs this loop as a serverless Databricks Job. The notebook at scripts/databricks/rag_evaluation.py queries five candidates from the real AI Search index, uses the configured Databricks model to retain only sources that supply relevant or complementary evidence, records the filtered documents in the retriever span, and generates a citation-backed answer. It merges the governed 50-case golden set into a content-hashed Unity Catalog evaluation dataset, so changed fixtures cannot inherit stale rows from an earlier release, and runs the following managed judges:

The source documents and cases live together in scripts/databricks/rag_fixture.json. The provisioner merges those documents into the Delta source table, triggers the AI Search index, and embeds the same cases into the evaluation notebook. This prevents the index fixture and golden set from drifting apart.

Each aggregate must meet DATABRICKS_RAG_EVAL_THRESHOLD (default 0.8). Evidence includes the MLflow run, versioned dataset and fixture hash, generation and judge models, category counts, selected aggregate metrics, and threshold result. A failed judge fails the Databricks Job and therefore the release certification gate.

The protected release workflow for commit 5fcf927bcf0a3a15e3aec86e2629516cbb76ad26 evaluated fixture 18bd4dffc32c as MLflow run 18131356bcac4460a26de20d40b9573a in the reference Azure workspace. It covered all 50 cases with databricks-gpt-oss-120b for generation and managed judging, and passed the unchanged 0.8 floor:

This is retained workspace evidence for the reference fixture, not a claim that every application or dataset will achieve the same scores. Replace the documents, question-level expected facts, and adversarial cases with your domain corpus before treating the gate as production evidence.

flowchart LR
  G[UC golden dataset<br/>content-hashed version] --> J[Serverless evaluation Job]
  J --> V[AI Search<br/>top five candidates]
  V --> R[Model-assisted reranker<br/>relevant sources only]
  R --> T[MLflow retriever span]
  T --> M[AI Gateway generation<br/>bounded context and citations]
  T --> E[Five MLflow managed judges]
  M --> E
  E -->|all scores at least 0.8| P[Release evidence]
  E -->|score below threshold| F[Block release]
  F --> D[Diagnose retrieval, prompt, or index]
  D --> J

  classDef source fill:#dcfce7,stroke:#16a34a,color:#052e16
  classDef gate fill:#fef3c7,stroke:#d97706,color:#422006
  classDef fail fill:#fee2e2,stroke:#dc2626,color:#450a0a
  class G,V,R,T,M source
  class J,E,P gate
  class F,D fail

Provision or repair the disposable fixture, then run the same gate used by protected CI:

FABRIC_DATABRICKS_PROVISION=1 pnpm databricks:cert:provision
pnpm databricks:certify

The provisioner is idempotent: it reuses the UC dataset, Jobs, AI Search index, Feature Serving endpoint, Genie Agent, and Lakeflow pipeline. Improve quality by diagnosing retrieval versus generation, changing topK, prompts, chunks, or embeddings, and rerunning the managed judges before redeploying Apps or Model Serving.

The default job limits MLflow to one prediction worker and one scorer worker, with prediction and scorer rate limits, so pay-per-token endpoints do not exceed workspace output-token quotas. Raise MLFLOW_GENAI_EVAL_MAX_WORKERS, MLFLOW_GENAI_EVAL_MAX_SCORER_WORKERS, MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT, or MLFLOW_GENAI_EVAL_SCORER_RATE_LIMIT only when the workspace uses capacity that supports the additional concurrency.

Do not build a second full judge product in Fabric when MLflow managed evaluation already exists in the workspace.

Offline index pipeline (not the chain)

Building/updating the AI Search index (chunking, embedding, write) stays on Databricks:

• Notebooks / Jobs
• Lakeflow pipelines (fh add lakeflow)
• UC Volumes as document sources

Fabric agents consume the index via AI Search at inference time.

Guardrails

The default chain treats retrieved chunks as untrusted data, serializes them as bounded JSONL, limits per-chunk and total context size, and rejects citation markers that do not match a retrieved source. turn.sources means "retrieved"; only source IDs referenced by the final answer appear in turn.citations.

const chain = databricksRagChain({
  databricks: { /* model + aiSearch */ },
  topK: 5,
  maxContextChars: 32_000,
  maxChunkChars: 8_000,
  postProcess: {
    citationPolicy: 'validate',
    maxAnswerChars: 8_000,
  },
});

See also

• Databricks recipes
• Integrations map
• example: with-databricks-rag
• Cookbook: RAG chain for inference