Skip to content

Library · Data Practice OS

AI on Databricks

Reference guide · Published · Source reviewed

Databricks' AI story is aimed at builders. You get the parts to build, evaluate, serve and monitor your own models and agents, all governed by Unity Catalog and tracked in MLflow. If the goal is a custom agent, a trained model or document processing at scale, this is the more complete toolkit.

Pick the component by the job

Job Component Notes
Classify, extract, summarise or parse documents in batch Databricks AI Functions (ai_query, ai_classify, ai_extract, ai_parse_document) LLM functions in SQL in pipelines; scales to large batch inference
"Ask the data" in plain English AI/BI Genie spaces over curated tables and metric views Text-to-SQL tuned with instructions and example SQL
Retrieval for RAG Databricks AI Search (formerly Vector Search) Delta-synced index with hybrid search and Unity Catalog permissions
Task-specific agents quickly (extraction, knowledge assistant, supervisor) Agent Bricks Declarative. Generates evaluation data and tunes quality against cost
Custom agents in code (LangGraph, OpenAI Agents SDK, plain Python) Agent Bricks (code-first agents) Tools are Unity Catalog functions; tracing via MLflow
Evaluate an agent or RAG app MLflow GenAI evaluation (Agent Evaluation) LLM judges plus your own labelled set. See Evaluating AI systems
Govern model access, rate limits and logging Model Serving with AI gateway One endpoint per approved model; usage and payload logging to tables
Classical ML Feature engineering in Unity Catalog, MLflow registry, Model Serving Feature store and Model registry
App front end and state Databricks Apps and Lakebase Agent memory and app state in Postgres, next to the lakehouse

Anthropic Claude models are available as Databricks-hosted foundation model endpoints. Verify which models are served in Sydney and whether any route via cross-geography processing in the documentation; see ADR-005: Governed model access through an AI gateway.

Keep it governed

  • Unity Catalog owns everything. Models, vector indexes, functions used as tools, and agents are UC objects with grants and lineage. An agent's tool can read only what its owner or caller can.
  • Route every model call through Unity Gateway (formerly AI Gateway): authentication, per-team rate limits, usage tracking, payload inference tables and Guardrails. Do not issue raw provider API keys to teams.
  • Turn off cross-geography processing for workspaces that handle SENSITIVE data, and record any exception. See Data residency.
  • Expose tools with managed Model Context Protocol servers (Unity Catalog functions, AI Search, Genie, Databricks SQL; Public Preview) so external assistants get the same governed access.

Patterns that work

1. Batch enrichment in the pipeline.

CREATE OR REPLACE TABLE safety.silver.crash_narrative_enriched AS
SELECT crash_id,
       narrative,
       ai_classify(narrative,
         ARRAY('speed', 'fatigue', 'distraction', 'alcohol_drugs', 'road_condition', 'other')
       ) AS primary_factor_ai
FROM safety.bronze.crash_reports
WHERE ingest_date = current_date();

2. Document-heavy RAG. Parse PDFs and scans with ai_parse_document, chunk them into a Delta table, and sync them to an AI Search index (formerly Vector Search). Filter by classification and business unit at query time.

3. Agent in code, evaluated before release. Build the agent with Agent Bricks, log it to MLflow, register it in Unity Catalog, run the evaluation set as a CI gate, then deploy it to a serving endpoint behind Unity Gateway. Production traces feed back into the evaluation set. Databricks can synthesise a first evaluation set from your documents, but you choose or write the MLflow judges and align them to expert grades. See Evaluating AI systems.

When not to use it

  • Business users who just need governed Q&A over a well-modelled SQL estate already in Snowflake. Cortex Analyst and Snowflake Intelligence get there with less engineering (AI on Snowflake - Cortex).
  • One-off experiments with no path to production. Use a sandbox with a budget and no access to SENSITIVE data.

Pitfalls

  • Notebook prototypes that never get an evaluation set, so nobody can say whether a change made them better.
  • Vector indexes over raw, unclassified documents. If PII is indexed, it can be retrieved.
  • Defaulting to the largest model. Agent Bricks and evaluation exist to find the cheapest model that meets the quality bar.
  • Leaving serving endpoints scaled up. Use scale-to-zero for low-traffic endpoints (Cost control on Databricks).

Worked example

The road asset inspection team captures thousands of defect photos and inspector notes a week. A Lakeflow pipeline lands the images and notes. A vision-capable model, called through ai_query, classifies defect type and severity. An Agent Bricks information-extraction agent pulls structured fields from contractor PDF reports. Results go to a gold asset_defect table with an _ai confidence column, and planners review low-confidence items (Human in the loop) before work orders are raised. MLflow evaluation, against 200 inspector-labelled examples, gates every prompt or model change.

Note on agent memory

Storing agent state in Lakebase is a storage choice, not a memory policy. Scopes, write rights, review and retention are an open decision in most organisations. See Agent memory before enabling long-term memory.