Divyesh Patel

July 9, 2026

5 Minutes read

Building Intelligent Agents on the Databricks Stack

The shift from single models to multi-step, reasoning agents introduces complexity, making governance, reliability, and security harder to manage at enterprise scale. Databricks and MLflow 3.0 provide the stack to solve this — from raw data to deployed, governed, production-grade agents.

What "Agentic AI" Actually Means

The term agentic AI has graduated from buzzword to engineering specification. An AI agent is any system where a language model iteratively calls tools, accumulates context, and makes multi-step decisions to complete a goal without constant human hand-holding between steps.

The dominant architectural pattern driving this is ReAct (Reason + Act): the model reasons about what to do next, takes an action (API call, SQL query, Python execution, search), observes the result, and reasons again. This loop continues until the task is complete or a stopping condition is met.

Agents differ from chains (fixed pipelines) in one critical way: the model decides the next step based on context, not a hard-coded graph. This makes them powerful and harder to govern.

The Four Pillars of an Agent

Planning & Reasoning: The LLM backbone decomposes goals into sub-tasks, evaluates options, and decides what tool to invoke next.
Tool Use: Structured function-calling lets agents execute SQL, call REST APIs, run Python, or query vector stores — extending far beyond text generation.
Memory: Short-term (in-context), episodic (summarised conversation), and long-term (vector/relational stores) memory gives agents continuity across sessions.
Reflection & Self-Correction: Advanced agents critique their own outputs, retry failed tool calls, and re-plan when they detect they have gone off-track.

The Databricks Lakehouse as Agent Infrastructure

Databricks was designed around one principle: unify data, analytics, and ML on a single platform. That design pays enormous dividends when you need to deploy agents that must reason over structured enterprise data.

Unity Catalog as Agent Context

Unity Catalog’s governed metadata layer isn’t just for governance — it’s a live directory your agents can query. An agent can discover what tables exist, understand their schema and lineage, and formulate appropriate SQL queries without human-in-the-loop schema lookup.

Mosaic AI Model Serving

Databricks’ dedicated inference layer supports streaming, autoscaling, and pay-per-token pricing for external models. Agents require low-latency, high-throughput inference to keep loop latency tolerable. Model Serving delivers exactly that, with governance baked in.

Vector Search as Agent Memory

The built-in Databricks Vector Search service lets agents do semantic retrieval over Delta tables. This bridges the gap between structured warehouse data and unstructured documents — critical for RAG-augmented agentic workflows.

MLflow 3.0 and the Agentic Lifecycle

MLflow, open-sourced by Databricks in 2018, was originally a lightweight experiment tracker. MLflow has evolved into a comprehensive lifecycle platform for agentic systems, covering tracing, evaluation, and deployment of multi-step agent systems.

Native Agent Tracing

MLflow’s tracing module instruments every LLM call, tool invocation, and retrieval step in an agent run. Each trace is a structured, queryable artefact that captures the full decision sequence: what the model reasoned, which tool it called, what the tool returned, and how that result influenced the next step. This makes debugging deterministic rather than speculative — engineers can replay any agent decision from any historical run, rather than trying to reproduce failures from logs alone.

In practice, this means wrapping the agent’s execution logic within a tracing context so that every internal step, planning cycles, tool dispatches, retrieval queries, and final synthesis is automatically captured. The resulting traces expose the latency and token cost associated with each granular step, enabling precise performance optimization and a complete audit trail of how the agent reached a specific conclusion.

Agent Evaluation with LLM-as-Judge

Evaluating agent quality has historically been a manual, expensive, and inconsistent process. MLflow’s evaluation framework addresses this by accepting agent traces as input and running configurable LLM judges against them, scoring for faithfulness, groundedness, toxicity, and task completion. This converts subjective assessment of agent quality into trackable, comparable metrics that can be monitored across model versions and prompt iterations.

The workflow connects the MLflow evaluation API to the agent’s deployment URI, enabling automated testing against gold-standard datasets. LLM judges systematically analyse traces to verify that the retrieved context is genuinely relevant to the query (groundedness) and that the final response directly addresses the user’s intent (relevance). Integrating these metrics into a CI/CD pipeline ensures that new model versions or prompt changes do not silently regress in safety or accuracy before reaching production users.

MLflow Model Registry for Agents

Agents are logged in MLFlow as standard PyFunc models, the same model format used for traditional ML. This means your promotion pipeline (staging → champion → production), A/B testing infrastructure, and rollback mechanisms work identically for LLM agents as for gradient-boosted trees.

MLflow Feature Comparison: Traditional ML vs Agentic AI

MLflow Feature	Traditional ML Use	Agentic AI Use
Tracking	Loss curves, accuracy	LLM call count, latency, token cost
Artifacts	Model weights, plots	Traces, system prompts, tool schemas
Evaluation	RMSE, F1, AUC	Groundedness, faithfulness, safety
Registry	sklearn / PyTorch models	Agent graphs as pyfunc models
Serving	REST endpoints	Streaming agent endpoints

Enterprise Agentic Pipeline

A production-grade Databricks agent is built around a modular design pattern with four conceptually distinct layers, each with clear responsibilities and governed interfaces. Before breaking these layers down individually, it helps to see how they connect end-to-end from raw data ingestion through to a monitored, served agent.

Figure 1: The six-stage agentic lakehouse pipeline, from data ingestion through governed monitoring, with a continuous feedback loop closing the cycle.

This pipeline reads left to right as a single governed lifecycle, with every stage backed by a specific Databricks or MLflow primitive rather than a generic placeholder:

Ingest: Delta Live Tables handle streaming and batch ingestion, change-data-capture from operational databases, and automatic schema evolution, so the agent’s data foundation stays current without manual pipeline maintenance.
Govern: Unity Catalog applies column-level security and lineage tracking to that ingested data, and exposes it as agent-readable metadata in the same governed layer discussed earlier as “Unity Catalog as Agent Context.”
Reason: The Agent Core executes the ReAct loop (Thought → Action → Observation) and dispatches tool calls, typically implemented through an orchestration framework such as LangGraph or the Databricks AI Agents SDK.
Track: MLflow Tracing captures every step of that reasoning loop, runs LLM-as-judge evaluation against it, and registers the resulting agent in the Model Registry.
Serve: Model Serving exposes the registered agent via autoscaled, streaming endpoints, with A/B traffic splitting available to safely compare agent versions.
Monitor: Lakehouse Monitoring closes the loop with drift detection on inputs, quality dashboards, and alert pipelines that feed directly back into stage one.

That last point is what makes this a cycle rather than a one-way pipeline: insights surfaced during monitoring – drifted inputs, degrading quality scores, alert patterns feed back into re-ingestion, reasoning improvements, and model updates, closing the loop shown at the bottom of the diagram. Every stage in this lifecycle sits on the same underlying platform primitives: Delta Lake, Unity Catalog, MLflow, Serving, and Monitoring, unified under the Databricks Lakehouse Platform.

Zooming into the “Reason” stage and the tool layer beneath it reveals the finer-grained control flow that actually executes on each user query, API event, or scheduled trigger:

Figure 2:Request-level architecture flow: a Planner LLM, Tool Router, Executor, and Reflector operate over a shared tool layer, with every step captured by MLflow.

While Figure 1 shows the platform lifecycle, this diagram illustrates what happens within the “Reason” and “Track” stages for a single incoming request. An incoming trigger a user query, an API event, or a scheduled job enters a four-part control loop:

Planner LLM – understands intent and breaks the task down into sub-tasks, mirroring the Planning & Reasoning pillar described earlier.
Tool Router – selects the best available tools for each sub-task and plans the execution order.
Executor – actually invokes the selected tools and collects their results.
Reflector – validates those results against the original intent, and reflects and refines before the loop either continues or terminates with a final response.

That control loop calls out to a shared tool layer: SQL or Genie for Databricks SQL access, Vector Search for knowledge retrieval, a sandboxed Python REPL for arbitrary execution, REST for external APIs and services, and Unity Catalog for governed data access. Every interaction with this tool layer is then captured by the same MLflow primitives shown along the bottom of the diagram: a trace store that captures each step and artefact, an evaluation suite running LLM-as-judge and quality checks, the Model Registry handling versioning and approval, and finally a serving endpoint that deploys, serves, monitors, and scales the agent the same Track and Serve stages from Figure 1, now seen at request granularity rather than platform granularity.

Human-in-the-loop (HITL) for high-risk actions

For actions that are irreversible or carry a significant risk of deleting records, sending external communications, or approving financial transactions, agents must pause and await explicit human approval before proceeding.

Pattern: pause-and-approve

Human-in-the-loop (HITL) is implemented as a special high-priority tool in the agent’s tool set. When the LLM determines that an intended action meets a pre-defined risk threshold, it invokes this tool rather than executing the action directly. The HITL tool triggers a state-persistence routine: the agent serialises its full conversation history and the pending action details to a secure, durable store, broadcasts a structured notification to designated human reviewers through enterprise communication channels, and enters a dormant state.

Execution resumes only when an authorised reviewer submits an approval or rejection through a secure callback. This design ensures that the agent’s operational efficiency is preserved for routine tasks, while human judgment is injected precisely at the boundaries where autonomous decisions carry unacceptable risk.

Token cost observability

Agentic loops multiply token usage unpredictably. A single agent run with ten tool calls and a long context window can cost 10–50x a single inference call. Without cost observability, budget overruns are invisible until the billing cycle closes.

Log costs as first-class MLflow metrics

To maintain financial governance, the agentic loop is instrumented with a cost-tracking module that captures token consumption for both input prompts and model completions. After each interaction with the LLM, the system calculates the estimated cost based on the specific model’s pricing tier. These financial metrics are then aggregated and logged to MLflow alongside performance data. This pattern provides immediate visibility into the economic impact of different query types and allows teams to set automated budget alerts, ensuring that the inherent unpredictability of agentic reasoning does not lead to unmanaged operational expenses.

Set MLflow alerts on estimated_cost_usd to catch runaway runs early.
Track cost per question type in your eval dataset to identify expensive query patterns.
Use token budgets: pass a remaining_token_budget value in the system prompt and instruct the model to be concise when it is low.

Security and Enterprise Governance

For agents to move from prototype to production, they must meet strict enterprise security and governance standards. The Databricks stack addresses this through three key areas:

Unified Governance via Unity Catalog: The Unity Catalog provides a single point of control for data, ensuring agents can only access authorized tables and metadata. This governed metadata layer is crucial for agents that formulate dynamic SQL queries, preventing unauthorized data access.
Execution and Credential Security: When agents execute code (like SQL or Python tools), a secure environment is non-negotiable. Enterprises must implement execution sandboxing (e.g., using secure containerization) for tools to mitigate risks associated with LLM-generated code. Furthermore, tools must use dedicated credential management systems (like secrets managers) instead of hard-coded keys to handle privileged access securely.
Policy and Auditability (Combating Shadow AI): Deploying agents requires strict adherence to corporate AI Usage Policy & Guideline and a governance framework to prevent the proliferation of unauthorized or Shadow AI tools. MLflow tracing provides the mandatory audit trail, allowing developers to replay every decision (LLM call, tool invocation) for review against policy, ensuring compliance and model safety.

What Goes Wrong (and How to Fix It)

Infinite Tool Loops: Agents can get stuck retrying failed tools. Always set a maximum iteration count and log the termination reason in MLflow. Treat hitting the limit as a failure metric to monitor.
Context Window Bloat: Long agent runs accumulate tool outputs in the context window. Implement a summarisation step every N turns, and log the compressed context as an MLflow artifact for auditability.
Non-Deterministic Evals: LLM judges are themselves stochastic. Run evaluations at temperature=0, with multiple judge models, and track variance. MLflow’s eval framework supports this natively in v3+.
Schema Drift Breaking Tool Calls: If the underlying Delta table schema changes, tool calls silently fail or hallucinate results. Add Unity Catalog schema-change alerts and route them to your agent’s CI/CD pipeline.
Treating Prompt as Config, Not Code: The system prompt is the most impactful part of your agent. Version it in MLflow, diff it between runs, and never deploy a prompt change without a paired evaluation run.

The Lakehouse Agent Stack: Closing Thoughts

Databricks gives you the data layer, compute, and inference infrastructure. MLflow gives you the reproducibility, governance, and deployment rails. Agentic AI gives you the paradigm shift. Together, they form the most complete production AI platform available today. Start small: instrument one existing pipeline with MLflow tracing, add a single tool call, and measure latency and quality. The architecture scales; your intuition needs to grow with it.

FAQs

What is an AI agent?

Unlike fixed AI chatbots that follow rigid scripts, an AI agent uses a language model to iteratively call tools, gather context, and make autonomous, multi-step decisions to complete a goal without constant human intervention.

How does the Databricks Lakehouse function as agent infrastructure?

It unifies data, security, and AI into a single platform. Unity Catalog acts as a live directory for agents to discover data schemas; Mosaic AI Model Serving provides low-latency inference; and Vector Search serves as the agent’s long-term memory across structured and unstructured data.

How do you prevent agents from executing unauthorized or high-risk actions?

By implementing a Human-in-the-Loop (HITL) pause-and-approve pattern. For irreversible actions (like approving financial transactions), the agent serializes its state, pauses execution, and alerts a human reviewer. It remains dormant until explicit approval is submitted.

How do you manage unpredictable agent token costs?

Instrument the agentic loop to calculate token spend per run and log the financial metrics directly to MLflow. You can then set automated budget alerts, track expensive query patterns, or inject a token budget directly into the prompt to force the agent to be concise.