Building an AI agent dbt data quality checks pipeline turns brittle manual tests into adaptive monitors that catch schema drift and statistical outliers before they hit production. This guide walks through a concrete implementation: a Python agent that introspects your dbt manifest, calls an LLM to propose validation rules, and executes them as dbt tests in CI.
Prerequisites
You need a dbt project (Core 1.6+), Python 3.11+, and an LLM API key. The agent emits standard dbt test definitions, so no custom runner is required. Install dependencies:
pip install dbt-core openai pyyaml pydanticSet your gateway key in the environment. We route the agent’s completions through n4n.ai, an OpenAI-compatible endpoint covering 240+ models with automatic fallback when a provider is rate-limited or degraded, so a single base_url swap works and you never hard-code a vendor.
export N4N_API_KEY="sk-..."
export N4N_BASE="https://api.n4n.ai/v1"Step 1: Capture dbt artifacts in CI
dbt emits target/manifest.json and target/run_results.json during parse and run. The agent needs the manifest to know models, columns, and data types. Run dbt parse before invoking the agent so the file is fresh:
dbt parse --target prod --output target/manifest.jsonIn GitHub Actions, persist it as an artifact:
- name: Parse dbt
run: dbt parse
- uses: actions/upload-artifact@v4
with:
name: dbt-manifest
path: target/manifest.jsonIf you want post-run stats (row counts, execution times), also keep target/run_results.json. The agent can use those to prioritize high-volume models.
Step 2: Extract model metadata for the agent
The manifest is large (often 5–20 MB). Trim it to the fields an LLM needs: model name, schema, column names, types, and any existing tests. This keeps token usage predictable.
import json
from pathlib import Path
def load_models(manifest_path: str) -> list[dict]:
manifest = json.loads(Path(manifest_path).read_text())
models = []
for uid, node in manifest["nodes"].items():
if node["resource_type"] != "model":
continue
cols = {
c["name"]: c.get("data_type", "unknown")
for c in node.get("columns", {}).values()
}
models.append({
"name": node["name"],
"schema": node["schema"],
"columns": cols,
"raw_sql": node.get("raw_code", "")[:2000], # truncate
"existing_tests": [
t["name"] for c in node.get("columns", {}).values()
for t in c.get("tests", [])
],
})
return models
models = load_models("target/manifest.json")
print(f"Loaded {len(models)} models")For a 50-model project this is roughly 8–12k tokens. Well within context for any modern model.
Step 3: Define the AI agent dbt data quality checks contract
We do not let the model free-write SQL. We constrain output to a JSON schema describing tests. The agent returns a list of proposed checks per model. Use pydantic to validate before writing files.
from pydantic import BaseModel, Field
from typing import Literal, Optional
class TestSpec(BaseModel):
model: str
column: Optional[str] = None
type: Literal["not_null", "unique", "accepted_values", "relationship", "custom_sql"]
params: dict = Field(default_factory=dict)
rationale: str
class AgentOutput(BaseModel):
tests: list[TestSpec]The custom_sql type maps to a dbt singular test. Everything else maps to built-in generic tests. This contract is the guardrail that keeps the AI agent dbt data quality checks output safe to apply.
Step 4: Implement the agent loop
We call an LLM with a system prompt enforcing the schema and a user prompt containing the model metadata. Because n4n.ai honors client routing directives and forwards provider cache-control hints, you can pin anthropic/claude-3.5-sonnet as primary and let it fall back to a similar model if that provider is degraded—no code change required.
from openai import OpenAI
import os, json
from pydantic import ValidationError
client = OpenAI(
base_url=os.environ["N4N_BASE"],
api_key=os.environ["N4N_API_KEY"],
)
SYSTEM = """You are a data quality agent. Given dbt model metadata, propose tests.
Return strict JSON matching the contract. Prefer built-in tests; use custom_sql only for cross-table or distribution checks."""
def propose_tests(model: dict) -> list[TestSpec]:
resp = client.chat.completions.create(
model="anthropic/claude-3.5-sonnet",
messages=[
{"role": "system", "content": SYSTEM},
{"role": "user", "content": json.dumps(model)},
],
response_format={"type": "json_object"},
temperature=0.2,
)
try:
return AgentOutput.model_validate_json(
resp.choices[0].message.content
).tests
except ValidationError as e:
print(f"Bad output for {model['name']}: {e}")
return []
all_tests: list[TestSpec] = []
for m in models:
all_tests.extend(propose_tests(m))The low temperature keeps output deterministic. Wrap the call in a tenacity retry if you run in flaky CI networks.
Step 5: Render dbt tests from agent output
Write the proposals into models/_ai_generated.yml. For generic tests use dbt’s YAML structure; for custom_sql emit a .sql singular test under tests/generated/.
import yaml
from pathlib import Path
def to_dbt_yaml(tests: list[TestSpec]) -> dict:
out = {"models": []}
for t in tests:
if t.type == "custom_sql":
continue
entry = next((m for m in out["models"] if m["name"] == t.model), None)
if not entry:
entry = {"name": t.model, "columns": []}
out["models"].append(entry)
col = next((c for c in entry["columns"] if c["name"] == t.column), None)
if not col:
col = {"name": t.column, "tests": []}
entry["columns"].append(col)
col["tests"].append({t.type: t.params})
return out
Path("models/_ai_generated.yml").write_text(
yaml.dump(to_dbt_yaml(all_tests))
)
for t in all_tests:
if t.type == "custom_sql":
sql = f"-- {t.rationale}\nselect * from {{{{ ref('{t.model}') }}}}\nwhere {t.params['where']}"
Path(f"tests/generated/{t.model}_{t.column}_check.sql").write_text(sql)After rendering, run dbt parse again to confirm the YAML is valid before test execution.
Step 6: Execute and verify in CI
Add a job that runs after the agent step. Start in report-only mode so the pipeline does not break on day one:
dbt test --select tag:ai_generated || trueTag the generated tests by adding tags: [ai_generated] to the YAML model entries if you want strict selection. After triage, flip to strict mode by removing || true.
How to verify success
A successful run meets three criteria:
models/_ai_generated.ymlexists anddbt parsevalidates it (no YAML errors).dbt testexecutes the new tests and returns exit code 0 on a known-good seed dataset.- The agent log shows proposed tests with rationales; spot-check that at least one test catches a seeded bad row.
To prove the agent works, inject a temporary bad row in a staging schema and confirm the test fails:
dbt seed --full-refresh && dbt test --select orders_amount_check
# expect non-zero exitRemove the seed override after verification.
Step 7: Schedule, diff, and human-review
Treat the AI agent dbt data quality checks output as a pull request, not a silent override. Run it nightly against the previous day’s manifest, open a PR with the YAML, and let a human approve. This keeps the LLM’s suggestions reviewable and stops drift from accumulating.
Cache the manifest between runs. Diff the manifest hash to skip unchanged nodes and cut token spend:
import hashlib
def manifest_hash(path: str) -> str:
return hashlib.sha256(Path(path).read_bytes()).hexdigest()[:8]If the hash matches the last run, reuse stored tests. If not, re-run the agent only on changed models.
Set a max token budget in the agent wrapper to fail closed if the manifest grows unexpectedly. Per-token metering on the gateway makes this observable per CI run.
Operational notes
Start with not_null and accepted_values only. Those are low-risk and build trust. Expand to custom_sql distribution checks once the team sees the agent catching real incidents. The goal is not full autonomy; it is a faster path from “we should test this” to “it’s tested”.
If a provider returns degraded latency, the gateway’s automatic fallback keeps the nightly job green. You do not need to babysit model routing.
The agent is a forcing function: it makes implicit assumptions explicit. Even when you reject half its proposals, you have a written rationale to argue against—better than a silent missing test.