AI Data Reliability Agent

Silent schema drift breaks ML models without a crash. This 5-agent LangGraph team catches it, traces DataHub lineage, and writes the incident back to DataHub.

Comment

Inspiration

Fraud detection models don’t crash when their data breaks. They just quietly get worse.

In FinTech, a model that scores every transaction often depends on just a handful of upstream PostgreSQL columns. When an engineer refactors an ETL pipeline and accidentally changes amount from DOUBLE PRECISION to VARCHAR, nothing necessarily throws an exception. The model keeps running. It just starts blocking legitimate customers and letting fraudulent transactions through. By the time anyone notices, the problem may have already affected business metrics for days.

We kept coming back to one observation: DataHub already knows the entire pipeline graph, including which tables feed which models, who owns them, and which columns are considered critical. Most tools treat that graph as a catalog that engineers browse. We wanted to see what would happen if an agent could actually act on it: read the graph to understand what needs protection, reason about what broke, and write the outcome back so the next engineer, or even the next agent, doesn’t have to start from a blank Slack thread.

What it does

We built a five-agent team consisting of Detector, RootCauseAnalyst, ImpactAssessor, FixAuthor, and Reviewer. They investigate a schema drift incident end to end, orchestrated as a real LangGraph state machine rather than a linear script dressed up as a multi-agent system.

Detector scans the live PostgreSQL schema against the baseline registered in DataHub and reads the audit trail.

RootCauseAnalyst walks the DataHub lineage graph upstream across multiple hops to identify every ML model that is actually affected, while also pulling ownership context.

ImpactAssessor and FixAuthor then run in parallel. ImpactAssessor calculates the real business impact using live row-level queries, while FixAuthor selects an appropriate remediation strategy from a vetted playbook based on the exact type of drift that occurred.

Reviewer statically validates the generated fix for unsafe SQL patterns and then asks an LLM for a final verdict based on those findings. The LLM can be stricter than the static validator, but it can never override a safety violation detected by the validator.

Only after the fix is approved does the agent write a first-class DataHub Incident back with an actual Active or Resolved status entity. This means anyone checking the affected model in DataHub can immediately see the warning instead of having to search through logs or Slack messages.

Every agent’s reasoning is genuine LLM output. None of the SQL that touches production is generated freely by the LLM.

How we built it

The system follows a “discovery, not configuration” approach.

At startup, the agent searches DataHub for production ML models based on their platform and tags. It then traverses the lineage graph upstream to discover the actual source tables. Nothing is hardcoded to the demo pipeline. Register a different model in DataHub, trigger Rediscover, and the agent automatically starts protecting it.

LangGraph compiles the five agents into a real execution graph. RootCauseAnalyst fans out into ImpactAssessor and FixAuthor, which run in parallel, before joining again at Reviewer. If the reviewer rejects a fix, the graph conditionally routes the workflow back to FixAuthor for another attempt. The review loop is bounded by a maximum iteration count so it cannot run indefinitely.

Tool access is also strictly separated between agents. Each agent receives a whitelist of tools it is allowed to call, and that restriction is enforced at the registry level. Detector cannot generate fixes, while FixAuthor cannot write to DataHub. This makes the agent handoff an actual system constraint rather than something that only exists in the prompt.

The remediation system supports three different strategies instead of relying on a single template.

When a column type changes, the system can restore the original type using an appropriate USING CAST operation.

When a column has been dropped, it cannot simply be cast back because the historical data is already gone. In that case, the system recreates the column and explicitly warns about the potential data loss.

When a completely new additive column appears, the column itself is not considered broken. Instead of rolling it back, the correct remediation is to register the new column into the baseline.

The LLM chooses and justifies the appropriate strategy, but the actual SQL comes from the corresponding vetted template.

The LLM layer is provider-agnostic. OpenAI is supported by default, including OpenAI-compatible gateways through a configurable base URL, with Anthropic also supported. There is no deterministic fallback pretending to be AI reasoning. If no LLM is configured, the system refuses to run rather than silently degrading into templated output.

The entire system also runs on real infrastructure. PostgreSQL 15 provides genuine information_schema introspection and an audit trail table. DataHub runs through its full Docker Compose reference stack with GMS, MySQL, Elasticsearch, and Kafka. A Streamlit control plane provides a live view of the agents, including their tool calls and reasoning as they happen.

Challenges we ran into

Most of our real challenges only appeared after we stopped testing against the happy path and started running the system against live infrastructure.

One of the biggest issues was DataHub’s datasetProperties behavior. The write operation is a full-aspect replacement rather than a merge. Our first incident-writing implementation accidentally replaced the model’s display name with the incident tag because we wrote the aspect without first reading the existing properties. The SDK did not warn us about this. We only discovered the problem by checking the actual DataHub UI and noticing that the model name had changed.

We also discovered that schema drift could be detected on the wrong node in the pipeline. A derived view between the raw table and the ML model was being treated as independently drifted because its inferred schema naturally differed from a separately registered baseline. We had to change the discovery process so that each source table is checked against its own upstream lineage, while excluding nodes that are not true raw sources from monitoring.

Another difficult issue came from PostgreSQL transaction handling. With autocommit=False, a single failed query can poison the entire session if the connection is not rolled back. For example, comparing a drifted VARCHAR column against an integer could fail and leave the connection in an aborted transaction state. Every subsequent query on that cached connection would then fail as well, including completely unrelated reads. This created a confusing symptom where a drift we had just introduced appeared to “disappear” during the next investigation. The root cause turned out to be a missing rollback().

Our DataHub lineage client also made an incorrect assumption about the GMS response structure. We expected relationships[].entity to contain a nested object, but the real DataHub GMS response returned it as a plain URN string. As a result, lineage calls were silently failing and falling back. We only caught this after deploying against a live GMS instance and inspecting the actual logs.

Finally, migrating the project to a new VPS exposed a Python version compatibility issue. An f-string containing an escaped quote inside its expression was valid on Python 3.12 and newer, but caused a hard SyntaxError on the older Python version provided by the Ubuntu environment. It worked perfectly during local development but caused the application to crash during startup in production.

Accomplishments we’re proud of

We achieved a genuine DataHub read/write round trip. The system reads six different aspects to build its protection map and writes three aspects back, including a first-class native Incident entity, with a documented fallback for older DataHub deployments.

We also built a safety invariant that we can actually defend under questioning. The Reviewer’s LLM verdict can be stricter than the static validator, but it can never be more lenient. If the static validator flags a script as unsafe, the system structurally prevents the LLM from approving it. This is enforced by code rather than simply being suggested through a prompt.

Another accomplishment was catching and fixing real bugs in our DataHub integration by looking at actual GMS responses and the actual DataHub UI instead of assuming the documentation perfectly matched reality.

Our remediation system also tells the truth about what it cannot fix. When a column has been dropped, the historical data is gone. The generated remediation explicitly communicates that limitation instead of pretending that a CAST can somehow recover deleted data.

What we learned

Metadata writes should always be treated as full-aspect replacements until proven otherwise. Reading the existing aspect before writing is the safest way to avoid silently overwriting fields that were never meant to change.

A lineage graph tells you what is connected, but not necessarily what is structurally equivalent. A raw table and a view built on top of it can both appear upstream, so the system needs additional logic to distinguish between actual source tables and derived nodes.

The most important boundary between “the LLM decides” and “the code executes” should exist around operations that can destroy data. Root-cause reasoning and business-impact analysis are areas where an LLM can be genuinely useful. Allowing an LLM to freely generate SQL that executes against production is a completely different risk.

Most importantly, bugs that never appear in a demo environment tend to surface as soon as the same code is connected to real infrastructure with slightly different behavior. Almost every major fix in this project came from actually running the system, inspecting the results, and tracing the failure back to its source rather than simply rereading the code.

What’s next for AI Data Reliability Agent

The next step is turning the local DataHub Skill specifications in datahub_skills/ into an actual upstream contribution to datahub-project/datahub instead of keeping them as repository-local artifacts.

We also want to build a “shift-left” mode that evaluates proposed schema migrations against the DataHub lineage graph before they are merged, rather than waiting until the migration has already broken something in production.

Another priority is validating the native Incident write-back across a wider range of real DataHub GMS versions. We also want to extend the same safety-gated approach to other types of schema drift, including column renames, constraint changes, and partition-key changes.

Finally, we plan to add notification integrations such as Slack and Microsoft Teams so that the Reviewer’s verdict can reach the team responsible for the affected model without requiring anyone to manually check a dashboard.

Built With

Share this project:

Updates