Agent Bricks: Building Enterprise AI Agents on Databricks

Everyone is building AI agents. Few are building agents that survive contact with enterprise reality — where data is governed, access is audited, and “it works on my laptop” doesn’t cut it.

Agent Bricks is Databricks’ answer to this gap. Instead of hand-wiring LLMs, retrieval pipelines, prompt templates, and evaluation harnesses, you describe what you want the agent to do, point it at your governed data, and the platform handles optimization, evaluation, and deployment.

In this article, we’ll go beyond the pitch deck: build a Knowledge Assistant, wire up a Supervisor Agent, deploy a custom agent, and understand the research (TAO, ALHF) that makes the auto-optimization pipeline work. With real code you can run today.

Why Agent Bricks Exists

#

Building a production AI agent on most platforms looks something like this:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

[Task Description]
      |
  Manual prompt engineering (days)
      |
  Manual evaluation suite (days)
      |
  Manual retrieval tuning (days)
      |
  Manual deployment pipeline (days)
      |
[Production Agent] (maybe)

That’s weeks of engineering before you know if the thing even works well enough to ship. And every time your data changes, you repeat parts of the cycle.

Agent Bricks collapses this into a four-step automated pipeline:

1
2
3
4
5
6
7
8
9

[Task Description + Data Sources]
      |
  Auto-generated evaluation benchmarks
      |
  TAO optimization sweep (prompts, models, retrieval, chunking)
      |
  Cost-quality curve → pick your tradeoff
      |
[Production Agent] (hours, not weeks)

The key insight: most of the engineering effort in agent development isn’t creative — it’s operational. Agent Bricks automates the operational parts while keeping you in control of the creative decisions.

Prerequisites

#

Before you start, make sure you have:

• A Databricks workspace with Unity Catalog enabled
• Serverless compute available in your workspace
• Access to Mosaic AI Model Serving
• The databricks-sdk Python package installed (pip install databricks-sdk)
• At least one data source ready (files in a UC Volume or a Vector Search Index)
• The databricks-gte-large-en embedding model available (required for Knowledge Assistants)

The Agent Bricks Lineup

#

Agent Bricks isn’t a single tool — it’s a family of agent templates, each targeting a common enterprise pattern:

The first three get the full auto-optimization treatment (TAO sweeps, synthetic benchmarks, LLM judges). Custom Agents trade that for complete architectural freedom — you bring your own framework, Databricks provides governance and deployment.

Let’s build each one.

Building a Knowledge Assistant

#

The Knowledge Assistant is the most common starting point. It’s a RAG-based agent that answers questions grounded in your enterprise data — but with a twist: Databricks’ “Instructed Retriever” approach claims 70% higher answer quality than naive RAG implementations.

Step 1: Create the Assistant

#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.agents import KnowledgeAssistant

w = WorkspaceClient()

assistant = KnowledgeAssistant(
    display_name="platform-docs-qa",
    description="Answers questions about our internal platform documentation",
    instructions=(
        "You are a helpful assistant for the platform engineering team. "
        "Answer questions based on the provided documentation. "
        "Always cite the source document. "
        "If you're not sure, say so — don't make things up."
    ),
)

created = w.knowledge_assistants.create(
    knowledge_assistant=assistant
)
print(f"Assistant ID: {created.name}")

Step 2: Connect Knowledge Sources

#

Knowledge Assistants support two types of data sources — UC Files and Vector Search Indexes. Let’s add both.

Adding files from a Unity Catalog Volume:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

from databricks.sdk.service.agents import KnowledgeSource, FilesSpec

files_source = KnowledgeSource(
    display_name="platform-runbooks",
    description="Operational runbooks and architecture decision records",
    source_type="files",
    files=FilesSpec(
        path="/Volumes/engineering/docs/runbooks"
    ),
)

w.knowledge_assistants.create_knowledge_source(
    parent=f"knowledge-assistants/{created.name}",
    knowledge_source=files_source,
)

Supported file types: txt, pdf, md, pptx, docx — up to 50 MB per file.

Adding a Vector Search Index:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from databricks.sdk.service.agents import KnowledgeSource, IndexSpec

index_source = KnowledgeSource(
    display_name="jira-tickets",
    description="Historical Jira tickets for incident context",
    source_type="index",
    index=IndexSpec(
        index_name="engineering.support.ticket_index",
        text_col="description",
        doc_uri_col="ticket_url",
    ),
)

w.knowledge_assistants.create_knowledge_source(
    parent=f"knowledge-assistants/{created.name}",
    knowledge_source=index_source,
)

Step 3: Let TAO Do Its Thing

#

Once your knowledge sources are connected, Agent Bricks kicks off the optimization pipeline automatically:

1. Synthetic benchmark generation — creates evaluation questions from your actual data
2. TAO sweep — searches across prompt strategies, chunking configurations, and retrieval parameters
3. Cost-quality curve — presents you with optimized configurations at different price points

You don’t write evaluation code. You don’t hand-tune chunk sizes. You pick a point on the cost-quality curve and deploy.

Step 4: Test in AI Playground

#

Before deploying, test your assistant in the Databricks AI Playground. Ask questions you know the answers to. Check that citations point to real documents. Try adversarial questions to see how the guardrails hold up.

Document Intelligence: Structured Extraction at Scale

#

If Knowledge Assistants are about answering questions, Document Intelligence is about extracting structured data from unstructured documents — contracts, invoices, clinical trial reports, compliance filings.

AstraZeneca used Document Intelligence to parse 400,000+ clinical trial documents in under 60 minutes, without writing extraction code.

The workflow is similar to Knowledge Assistants:

1. Define the extraction schema (what fields you want)
2. Point at your documents
3. Let the optimization pipeline tune the extraction
4. Deploy as a batch or real-time endpoint

The key difference: instead of answering free-form questions, Document Intelligence outputs structured records that can feed directly into your Delta Lake tables.

Supervisor Agent: Multi-Agent Orchestration

#

Here’s where it gets interesting. A Supervisor Agent coordinates up to 20 sub-agents, routing requests to the right specialist based on the task.

Think of it as a dispatcher:

1
2
3
4
5
6
7
8

[User Query]
      |
  Supervisor Agent
      |
  ┌───┼───┬───┬───┐
  v   v   v   v   v
 KA1 KA2 Genie UC  MCP
           Space Fn Server

Sub-agents can be:

• Knowledge Assistants (RAG over different data domains)
• Genie Spaces (natural language SQL over structured data)
• Unity Catalog Functions (deterministic business logic)
• External MCP Servers (third-party tool integration)

The Supervisor handles task delegation, result synthesis, and — critically — access control. End users only see sub-agents they’re authorized to use.

1
2
3
4
5

# Supervisor Agent configuration (via UI or API)
# Each sub-agent is registered with:
# - Display name and description (used for routing)
# - Access control (which users/groups can access it)
# - Timeout and retry configuration

Custom Agents: Bring Your Own Framework

#

Not every agent fits a template. For novel architectures, you can deploy agents built with any framework — LangChain, LangGraph, CrewAI, OpenAI Agents SDK, LlamaIndex, or pure Python — through Agent Bricks’ governed infrastructure.

The catch: Custom Agents don’t get TAO auto-optimization. You handle your own prompts and evaluation. But you still get Unity Catalog governance, MLflow tracing, and managed deployment.

The bridge is MLflow’s ResponsesAgent interface:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

import mlflow
from openai import OpenAI

class MyCustomAgent(mlflow.pyfunc.PythonModel):
    """Wrap any agent framework with the ResponsesAgent interface
    for compatibility with AI Playground, Agent Evaluation,
    and Databricks Apps."""

    def predict(self, context, model_input):
        # Your agent logic here — LangChain, CrewAI, raw API calls
        client = OpenAI()
        response = client.chat.completions.create(
            model="databricks-meta-llama-3-3-70b-instruct",
            messages=model_input["messages"],
        )
        return response.choices[0].message.content

# Log and register
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model=MyCustomAgent(),
        registered_model_name="my-custom-agent",
    )

Deploy to Databricks Apps:

1
2

databricks apps deploy my-custom-agent \
  --source-code-path /Workspace/Users/$USER/my-custom-agent

Once deployed, your custom agent gets the same governed infrastructure as built-in templates: Unity Catalog access control, MLflow experiment tracking, and Model Serving endpoints.

Under the Hood: TAO and ALHF

#

Two research innovations power the auto-optimization pipeline. Understanding them helps you reason about when Agent Bricks will work well — and when it won’t.

TAO (Test-time Adaptive Optimization)

#

Traditional model tuning requires labeled training data — expensive to create and quick to go stale. TAO flips this:

1. Takes unlabeled usage data (real queries your users would ask)
2. Uses test-time compute + reinforcement learning to optimize model behavior
3. Searches across prompts, retrieval configs, chunking strategies, and model choices simultaneously

The result: open-source models (like Llama) tuned with TAO can match GPT-4o quality on domain-specific tasks, at a fraction of the cost.

1
2
3
4
5
6
7
8

Quality
  ^
  |        ● GPT-4o (baseline)
  |      ● TAO-tuned Llama 3.3 70B
  |    ● Default Llama 3.3 70B
  |  ● Naive RAG
  |
  +──────────────────────> Cost

TAO scales with compute budget, not human labeling effort. More GPU hours = better optimization, without anyone writing a single evaluation example.

ALHF (Agent Learning from Human Feedback)

#

RLHF (Reinforcement Learning from Human Feedback) changed LLM training, but it has a limitation: feedback is binary (thumbs up/down). ALHF extends this with natural-language corrections.

Instead of rating an agent response good/bad, domain experts write what was wrong:

“The agent cited the 2024 policy, but the 2025 revision supersedes it. It should prioritize documents by recency for compliance questions.”

ALHF translates this into technical adjustments:

• Retrieval algorithm changes (recency weighting)
• Prompt modifications (prioritize recent sources)
• Vector database filtering updates
• Agentic pattern changes

This avoids the brittleness of cramming every edge case into a single system prompt. The corrections become structural improvements, not prompt patches.

Agent Bricks vs. DIY Frameworks

#

The honest comparison:

When to use Agent Bricks

#

• Your data is already in Unity Catalog
• You need governance and audit trails (regulated industries)
• You want to skip the evaluation/optimization engineering
• Your use case fits a template (knowledge Q&A, document extraction, multi-agent orchestration)

When to reach for LangChain / CrewAI

#

• You need architectural patterns Agent Bricks doesn’t support
• Multi-cloud portability is a hard requirement
• You want to own every component of the stack
• Your agent needs novel tool-use patterns that don’t fit the Supervisor model

The hybrid path

#

Use Custom Agents to deploy your LangChain/CrewAI agent through Agent Bricks infrastructure. You lose auto-optimization but keep governance, deployment, and monitoring. This is often the pragmatic middle ground.

Production Customers

#

Agent Bricks isn’t just a demo — it’s running at scale:

• AstraZeneca — 400,000+ clinical trial documents processed with Document Intelligence
• Workday — Employee-facing knowledge assistants across HR documentation
• Virgin Atlantic — Customer service agents grounded in operational data
• Zapier — Internal tooling agents for workflow automation
• EchoStar — Multi-agent systems for satellite operations

Getting Started: Your First 30 Minutes

#

1. Enable prerequisites — Unity Catalog, serverless compute, Model Serving
2. Create a Knowledge Assistant in the Databricks UI (Mosaic AI > Agent Bricks)
3. Add a knowledge source — start with a small UC Volume of markdown or PDF files
4. Wait for the optimization sweep — this runs automatically
5. Test in AI Playground — ask questions, check citations, try edge cases
6. Deploy — pick a point on the cost-quality curve and ship it

You’ll have a governed, production-ready RAG agent in under an hour. Whether that’s impressive or terrifying depends on how long your last agent project took.

Limitations and Gotchas

#

Before you go all-in:

• English only — no multilingual support yet
• File size limit — 50 MB per file for Knowledge Assistants
• Embedding lock-in — Vector Search Indexes must use databricks-gte-large-en
• Supervisor cap — maximum 20 sub-agents
• No TAO for Custom Agents — auto-optimization only works on built-in templates
• Enhanced Security workspaces — Supervisor Agent has compatibility limitations
• Platform dependency — Agent Bricks is tightly coupled to the Databricks ecosystem. If you’re multi-cloud or considering platform migration, factor this in

What’s Next

#

Agent Bricks represents a bet that most enterprise agent development is more operational than creative. If that’s true for your use case, it eliminates weeks of engineering. If your agent needs architectural novelty, the Custom Agents path gives you governance without constraints.

Either way, the direction is clear: the platform is eating the framework. Databricks isn’t just providing infrastructure for AI agents — it’s automating the development process itself.

The question isn’t whether your enterprise needs AI agents. It’s whether you’ll spend weeks building the plumbing, or let the platform handle it while you focus on what the agent actually needs to do.

Agent Bricks is Generally Available on Databricks. Documentation: AWS | Azure