Prognosis — Self-Healing Data Agent for DataHub

Your data's immune system. Traces silent ML drift to its root cause in DataHub's lineage, opens a real dbt fix PR, and vaccinates the graph so the next agent inherits the knowledge.

Comment

Inspiration

A production recommendation model loses 12% accuracy over three days. Nobody changed the model. Nobody changed the training code. Upstream, in a source table three hops away, someone renamed a column from user_region to user_geo — and the whole pipeline broke without a single alert.

This failure mode is brutally common and brutally expensive. The pipelines keep running. The dashboards keep rendering. The model keeps predicting — it just predicts worse. Teams burn days investigating the model before discovering the cause was never in the model at all.

The information needed to diagnose this in seconds already exists: it's sitting in DataHub's lineage graph. Nobody queries it at the right moment, and nothing writes the answer back once it's found. So the next team, six months later, re-lives the same incident from scratch.

Prognosis closes that loop.

What it does

Prognosis is a self-healing data agent that treats DataHub's lineage graph as an immune system. Five specialized agents, orchestrated with LangGraph over a single shared MCP connection:

Reactive path — when a production model drifts:

  1. Triage — reads the entity from the graph (type, owner, production status) and rates severity. The owner is a fact pulled from DataHub; only the severity judgement comes from the LLM.
  2. Root-Cause — walks the lineage upstream hop by hop, reads every schema along the way via list_schema_fields, and identifies the exact break: a column propagated down the chain that has vanished at the source.
  3. Remediation — generates a dbt patch grounded in the real source schema (user_geo AS user_region) and opens an actual GitHub Pull Request.
  4. Immunization — writes the incident back into DataHub: a prognosis-incident tag, four structured properties, an enriched description, and a linked post-mortem document. This is the scar.

Preventive path — before the break ever happens:

  1. Pre-Flight — given a proposed schema change, it queries downstream lineage and computes a deterministic risk score, flagging every production ML model at risk. It posts a risk comment on the PR before merge. No simulation engine — just the lineage DataHub already has.

The result is a genuine closed loop: read → reason → act → write back. An incident solved once protects every future agent that touches that data.

How we built it

  • DataHub MCP Server (mcp-server-datahub) as the single interface to the graph — both read tools (search, get_entities, get_lineage, get_lineage_paths_between, list_schema_fields, get_dataset_queries) and mutation tools (add_tags, add_structured_properties, update_description, save_document).
  • LangGraph for orchestration, with one MCP connection injected across all nodes via functools.partial — critical for performance.
  • Gemini 3.6 Flash for reasoning only, never for facts.
  • PyGithub for real PR creation, with graceful fallback when no token is present (so judges can run it without one).
  • A real-time dashboard (FastAPI + WebSocket) that makes the agents' work legible as it happens — including a lineage graph that animates hop by hop, driven by actual MCP calls rather than a scripted animation.
  • A 22-point health check (python scripts/health_check.py) that verifies the entire system end-to-end with factual assertions on the pipeline output.

The governing design principle throughout: facts come from MCP, judgment comes from the LLM. The root-cause column is detected by deterministic code comparing real schemas; the LLM only explains why it matters. The generated SQL uses column names read from the live schema, never invented.

Challenges we ran into

ML lineage is asymmetric. Walking downstream from a feature table does not surface the ML model when that model is connected through a dataProcessInstance training run — the traversal stops at the dashboard. We solved it with an inverse cross-check: for each ML model, query its upstream lineage and check whether any threatened dataset appears there. Without this, Pre-Flight would have reported "0 ML models impacted" on an incident that threatens a production model. We documented this as a reusable pattern for the community.

Silent failures in the mutation layer. update_description returned a Pydantic validation error as a string rather than raising an exception, so our code logged "success" while nothing was written to DataHub. We only caught it by checking the UI rather than trusting the logs. The fix: parse and verify the tool's actual response, never infer success from the absence of an exception.

Naive schema diffing produces false positives. Our first root-cause implementation flagged any column present downstream but missing upstream — which wrongly accused computed columns like engagement_score. The correct signal is narrower: a column propagated across multiple consecutive datasets that then disappears at the source.

Free-tier rate limits. 5 requests/minute forced us to build proper LLM throttling — which turned out to be a feature, since it makes the agent's real thinking time visible rather than hidden.

Accomplishments that we're proud of

  • The loop genuinely closes. The scar is visible in DataHub's UI — tag, structured properties, documentation, post-mortem — not just claimed in a log.
  • The PR is real, cliquable, and merge-ready. Not a simulated URL.
  • Nothing is hallucinated. Every fact on screen traces back to an MCP call.
  • The dashboard's lineage animation runs at the speed of actual MCP calls, which is why the hops are seconds apart. It's an honest visualization, not a demo effect.

What we learned

Never trust a success log — verify the artifact. Three separate bugs in this project were hidden behind cheerful log lines. Checking DataHub's UI and reading raw tool responses caught what the logs concealed.

We also learned that the boundary between deterministic code and LLM reasoning is where agent reliability lives. Every time we let the model infer something the graph already knew, quality dropped. Every time we pushed a fact back into code, it improved.

What's next for Prognosis

  • Ship the RFC as a real PR: an atomic create_incident_from_lineage MCP tool collapsing the four-call immunization sequence into one.
  • Wire Pre-Flight to a live GitHub webhook so it runs automatically on every schema-touching PR.
  • Broaden detection beyond schema breaks — type changes, nullability shifts, freshness anomalies.
  • Multi-incident memory: query the accumulated scars to surface which datasets break most often.

Built With

Share this project:

Updates