Data quality checks are one of those tasks that every data team says they do-but almost no one does consistently. The usual workflow is: open a notebook, run a few SQL queries, eyeball the results, close the notebook, and forget about it until something breaks downstream.
We wanted to fix that. Not with a dashboard, not with a scheduled job, but with something simpler-an agent you can talk to. Type check data quality, get a full health report. No SQL, no notebooks, no waiting for an engineer.
This post documents exactly how we built that agent on Databricks using MLflow, the errors we hit (there were many), and what we learned from each one. By the end, you will have a clear understanding of what a custom agent actually is-not the buzzword version, but the real one.
1. What Is a Custom Agent?
Let us start here because the term gets thrown around loosely. A custom agent is not a large language model. It is not ChatGPT with your data. It does not reason or think.
A custom agent is your own code-your queries, your rules, your business logic-wrapped in a framework that lets anyone talk to it in plain English and get a structured, accurate response back.
The key distinction is this: the intelligence is yours. You write the predict() method. You decide what it checks. You control what it returns. The agent framework provides the plumbing-a way to register your code as a model, deploy it as an API endpoint, and give it a chat interface.
In Databricks, a custom agent is built using MLflow’s PythonModel class, registered in Unity Catalog as a versioned model, and deployed as a Model Serving Endpoint. Once live, it accepts messages like any chat API and returns responses in OpenAI-compatible format.
2. Why We Built This
Our team works with a Delta Lake table stored in Unity Catalog as workshops_db.agent_schema.ec_orders. It has 16 records, 10 columns, and a handful of known quality issues we wanted to surface quickly.
The traditional approach would be a notebook with SQL queries checking for NULLs, duplicates, and out-of-range values. That notebook would get stale, forgotten, and eventually deleted. We wanted something living and accessible-something any team member could query without touching a notebook.
The goal: deploy an agent that accepts natural language and returns a real, accurate data quality report from the live table.
3. The Dataset
The table is stored in Unity Catalog as workshops_db.agent_schema.ec_orders. Below is a breakdown of its structure and the quality issues we found.
Column | Type | Quality Notes |
order_id | STRING | Clean-Primary Key (15 unique values) |
customer_id | STRING | 1 NULL (6.2% of rows) |
product_id | STRING | Clean-15 unique values |
category | STRING | Clean-3 distinct categories |
region | STRING | Clean-4 regions |
amount | DOUBLE | 2 NULLs (12.5% of rows) |
quantity | INTEGER | Clean-avg 2.81 |
order_date | DATE | Clean-2024-01-01 to 2024-01-07 |
status | STRING | Clean |
created_at | TIMESTAMP | Clean |
The table also contains one duplicate record, bringing its overall health score to 97.5 out of 100-not bad, but not perfect. That is exactly the kind of nuance an agent should surface.
4. Architecture Overview
Before writing any code, it helps to understand how a Databricks custom agent works end-to-end. The lifecycle splits into two distinct phases.
4.1 Log Time-What Happens in the Notebook
This is where all the heavy computation lives. The notebook has full Spark access and can query Delta tables. We run all data quality checks against workshops_db.agent_schema.ec_orders and produce a formatted report string. That report gets baked into a PythonModel class, then logged to MLflow and registered in Unity Catalog.
Delta Table (workshops_db.agent_schema.ec_orders)
|
v
Spark Queries (NULL checks, duplicate detection, numeric stats)
|
v
PythonModel (report string baked in at log time)
|
v
MLflow Model Registry --> Unity Catalog: workshops_db.agent_schema.data_quality_agent
4.2 Serve Time-What Happens at the Endpoint
When someone sends a message to the deployed endpoint, the model’s predict() method runs. It parses the incoming message, matches it against keywords, and returns the pre-computed report. No database calls happen at serve time-the report is already in memory inside the model object.
User types: "check data quality"
|
v
Serving Endpoint (lightweight container, no JVM, no Spark)
|
v
predict() parses the message, matches keywords
|
v
Returns pre-computed report instantly
4.3 Why Pre-Compute Instead of Query at Serve Time?
This separation is the most important architectural decision in the project. Our first instinct was to run Spark queries inside predict() so each prompt would fetch fresh data. This failed immediately:
[JAVA_GATEWAY_EXITED] Java gateway process exited before sending its port number
Databricks serving endpoints run in lightweight containers. Those containers do not have a JVM. Spark requires a JVM. The solution is to move all Spark queries to log time inside the notebook and bake the resulting report into the model object. This is the pre-compute pattern.
5. Building It Step by Step
Step 1-Install Dependencies
Run the following cell at the top of your Databricks notebook. The library restart is required for the new packages to be available in the session.
%pip install databricks-agents mlflow pyspark
dbutils.library.restartPython()
Step 2-Pre-Compute Metrics Using Spark
This cell runs inside the notebook where Spark is available. We query workshops_db.agent_schema.ec_orders, compute every quality metric, and build the formatted report string that will be baked into the model.
from pyspark.sql import functions as F
from pyspark.sql.types import NumericType
TABLE = "workshops_db.agent_schema.ec_orders"
df = spark.sql(f"SELECT * FROM {TABLE}")
total_rows = df.count()
total_cols = len(df.columns)
duplicate_count = total_rows - df.dropDuplicates().count()
# NULL counts across every column
null_exprs = [F.sum(F.col(c).isNull().cast('int')).alias(c) for c in df.columns]
null_counts = df.select(null_exprs).collect()[0].asDict()
null_issues = [
f" {col}: {cnt} NULLs ({round(cnt / total_rows * 100, 1)}%)"
for col, cnt in null_counts.items() if cnt > 0
]
# Numeric column statistics
numeric_cols = [f.name for f in df.schema.fields if isinstance(f.dataType, NumericType)]
numeric_stats = []
if numeric_cols:
stats_df = df.select(numeric_cols).summary('min', 'max', 'mean')
stats = {row['summary']: row.asDict() for row in stats_df.collect()}
for col in numeric_cols:
mn = round(float(stats['min'][col]), 2)
mx = round(float(stats['max'][col]), 2)
avg = round(float(stats['mean'][col]), 2)
numeric_stats.append(f" {col} -> min: {mn}, max: {mx}, avg: {avg}")
# Health score calculation
total_nulls = sum(null_counts.values())
null_pct = round(total_nulls / (total_rows * total_cols) * 100, 2)
health_score = round(100 - null_pct - (duplicate_count / total_rows * 10), 1)
health_score = max(0, min(100, health_score))
health_label = 'GOOD' if health_score >= 80 else 'AVERAGE' if health_score >= 50 else 'POOR'
PRECOMPUTED_REPORT = '\n'.join(lines)
Step 3-Write the Agent Class
The DataQualityAgent class wraps the pre-computed report in an MLflow-compatible model. The predict() method handles both input formats that Databricks endpoints may send: a Pandas DataFrame or a plain dictionary.
from mlflow.pyfunc import PythonModel
import pandas as pd, json, mlflow
class DataQualityAgent(PythonModel):
def __init__(self, report):
self.report = report # baked in at log time, never re-queried
def predict(self, context, model_input):
try:
if isinstance(model_input, pd.DataFrame):
messages = model_input['messages'].iloc[0]
if isinstance(messages, str):
messages = json.loads(messages)
user_query = messages[-1]['content'].lower()
elif isinstance(model_input, dict):
user_query = model_input['messages'][-1]['content'].lower()
else:
user_query = ''
except:
user_query = ''
if any(w in user_query for w in ['quality', 'check', 'report', 'metrics']):
content = self.report
else:
content = 'Ask: check data quality'
return {
'choices': [{
'message': {'role': 'assistant', 'content': content},
'finish_reason': 'stop',
'index': 0
}]
}
mlflow.models.set_model(DataQualityAgent(PRECOMPUTED_REPORT))
Step 4-Log and Register the Model
We log the model to MLflow and register it in Unity Catalog under workshops_db.agent_schema.data_quality_agent. The signature is inferred using a plain string output-this avoids the deprecated rag_signatures schema that causes a ValueError on newer MLflow versions.
import mlflow
mlflow.set_registry_uri('databricks-uc')
input_example = {
'messages': [{'role': 'user', 'content': 'check data quality'}]
}
signature = mlflow.models.infer_signature(
model_input=input_example,
model_output="Data Quality Report-workshops_db.agent_schema.ec_orders"
)
with mlflow.start_run():
mlflow.pyfunc.log_model(
name="agent",
python_model=DataQualityAgent(PRECOMPUTED_REPORT),
registered_model_name="workshops_db.agent_schema.data_quality_agent",
input_example=input_example,
signature=signature
)
Step 5-Deploy to a Serving Endpoint
The deploy call registers a serving endpoint against the latest version of workshops_db.agent_schema.data_quality_agent. Wait five to ten minutes for the endpoint status to reach Ready before testing.
from databricks import agents
client = mlflow.MlflowClient()
versions = client.search_model_versions(
"name='workshops_db.agent_schema.data_quality_agent'"
)
latest_version = max([int(v.version) for v in versions])
agents.deploy(
model_name="workshops_db.agent_schema.data_quality_agent",
model_version=str(latest_version)
)
6. Every Error We Hit
We did not get this right in one shot. Here are the six errors we encountered, what caused each one, and exactly how we fixed it.
Error 1-Invalid Input Format
Error: Invalid input format returned immediately on first test
Databricks serving endpoints send model_input as a Pandas DataFrame, not a plain Python dictionary. Our isinstance(model_input, dict) check ran first and failed silently before it could do anything useful.
Fix: Add the DataFrame check before the dict check. Both paths must extract the last message’s content string and convert it to lowercase.
if isinstance(model_input, pd.DataFrame):
user_query = model_input['messages'].iloc[0][-1]['content'].lower()
elif isinstance(model_input, dict):
user_query = model_input['messages'][-1]['content'].lower()
Error 2-Wrong Response Format
Error: Playground threw a schema error after fixing the input
We were returning [{“content”: “…”}]. The Playground expects the OpenAI-compatible format: a choices array containing a message object with both role and content fields, plus finish_reason and index.
Fix: Return the full OpenAI-compatible choices structure on every predict() call, without exception.
return {
'choices': [{
'message': {'role': 'assistant', 'content': content},
'finish_reason': 'stop',
'index': 0
}]
}
Error 3-Spark Cannot Run Inside a Serving Endpoint
Error: JAVA_GATEWAY_EXITED-Java gateway process exited before sending its port number
The most significant error. Our original design ran Spark queries inside predict() so each prompt would fetch fresh data from the table. Databricks serving endpoints run in lightweight containers that have no JVM. Spark requires a JVM. The two are fundamentally incompatible.
Fix: Move all Spark queries to log time. Pre-compute the report in the notebook and bake it into the model object. The endpoint never touches a database or runs Spark.
Error 4-Deprecated MLflow Signatures
Error: ValueError: output schema not compatible with Agent Framework
A wall of deprecation warnings about mlflow.models.rag_signatures appeared, followed by a ValueError. Type hints on the predict() method were triggering schema validation logic that the agent framework could not parse.
Fix: Remove all type hints from predict(self, context, model_input). Use mlflow.models.infer_signature() with a plain string as the output example, not a typed schema object.
Error 5-SQL Connector Not Available in Container
Error: cannot import name ‘sql’ from ‘databricks’
After the Spark failure, we tried using databricks-sql-connector as an alternative way to query the table at serve time. The connector is not installed in the serving container and cannot be added.
Fix: Confirmed: the pre-compute pattern is the only correct approach. All data access must happen in the notebook at log time. There is no supported way to query Unity Catalog tables from inside a serving endpoint.
Error 6-Endpoint Currently Updating
Error: ValueError: endpoint is currently updating
After iterating through multiple fixes and redeploying quickly, we hit a conflict error. You cannot deploy a new version while a previous deployment is still in progress, even if that deployment failed.
Fix: Wait five to ten minutes between deployments. Check the endpoint status on the Serving page before calling agents.deploy() again.
7. What the Agent Actually Returns
Here is the exact output the agent produces when you type check data quality in the Playground. Every number is real-computed from actual Spark queries run against workshops_db.agent_schema.ec_orders at the time the model was logged.
Data Quality Report-workshops_db.agent_schema.ec_orders
=======================================================
Overview
Total Rows : 16
Total Columns : 10
Duplicates : 1
NULL Analysis
customer_id : 1 NULL (6.2%)
amount : 2 NULLs (12.5%)
Numeric Column Statistics
amount -> min: 100.0, max: 999999.0, avg: 149757.0
quantity -> min: 1.0, max: 10.0, avg: 2.81
Date Column Ranges
order_date -> from: 2024-01-01 to: 2024-01-07
String Column Uniqueness
order_id -> 15 unique values
customer_id -> 13 unique values
product_id -> 15 unique values
category -> 3 unique values
region -> 4 unique values
Overall Health Score : 97.5 / 100 (GOOD)
=======================================================
8. Supported Prompts
The agent recognises a range of natural language inputs. The keyword matching in predict() covers the following:
Prompt | What It Returns |
check data quality | Full quality report for the default table |
run quality report | Full quality report |
analyze the table | Full quality report |
check ec_orders quality | Report for that specific named table |
list tables | All available tables in workshops_db.agent_schema |
show null issues | Only the NULL analysis section |
show health score | Health score and rating only |
how many rows | Row and column count only |
show numeric stats | Min, max, avg for numeric columns |
Multi-Table Support
We extended the agent to automatically discover every table in workshops_db.agent_schema and pre-compute a report for each one. When a prompt contains a recognisable table name, the agent returns that table’s report. Adding a new table is simple – upload it to the schema, rerun the notebook cells, and redeploy. The agent picks it up automatically.
9. So – What Is a Custom Agent?
We started this post with a definition. After building through six errors and a full deployment, we can say it more concretely.
A custom agent is not magic. It is your code, running in a container, responding to messages. The “custom” part means you wrote the logic. The “agent” part means it has a chat interface and is accessible via an API.
In our case, the logic is five things:
- Query workshops_db.agent_schema.ec_orders with Spark
- Count NULLs, duplicates, and compute numeric column statistics
- Build a formatted report string
- Match user messages to keywords in predict()
- Return the report in the correct OpenAI-compatible JSON format
None of that is AI. All of it is Python. What the Databricks Agent Framework adds is a deployment target, a versioned model registry, a Playground UI for testing, and an API endpoint for integration.
A custom agent is the bridge between your data systems and the people who need answers from them-without requiring those people to know SQL, Spark, or Python. |
That is what makes it worth building. Not the technology, but the access it creates.
Our data quality agent now lives as a deployed endpoint (workshops_db.agent_schema.data_quality_agent) that any team member can query. No notebooks, no SQL, no waiting for an engineer. Just a question and an answer.
Key takeaways for anyone building something similar:
- Pre-compute all data at log time in the notebook-never query inside predict()
- Always return the OpenAI-compatible choices format from predict()
- Remove type hints from predict(self, context, model_input)
- Handle both Pandas DataFrame and dict inputs in your predict() method
Wait five to ten minutes between deployments-do not redeploy while an update is in progress