DataPilot AI : The Metadata Aware AI Data Engineer

An AI-powered data engineer that understands metadata, diagnoses data issues, and automates reliable data workflows.

Comment

DataPilot AI: The Metadata-Aware AI Data Engineer

Inspiration

Every data team has had the same experience: ask an AI coding assistant to "write a pipeline for our customer revenue table," and it hands back SQL that references customer_name, email, and country columns that don't exist, invented lineage that doesn't exist, and zero awareness that half the fields it just touched are tagged PII. The code looks production-ready. It isn't. It's a hallucination with good syntax.

We'd all been burned by this, and around the same time we were digging into DataHub's MCP server and realized the fix was obvious: an LLM doesn't need to guess an organization's schema if it can just ask. DataHub already knows the real columns, the real lineage, the real owners, the real governance tags. Nobody had wired that knowledge directly into the code-generation loop in a way that made hallucination structurally difficult rather than just "prompted against." That gap is DataPilot AI.

What it does

DataPilot AI is an agent that turns a plain-English request like "create a production pipeline that transforms raw.customer_transactions into a daily customer revenue model"into a real, validated, metadata-grounded pull request.

Concretely, it:

  1. Parses the request to identify the source dataset.
  2. Queries DataHub (via its official MCP server) for that dataset's real schema, column types, ownership, PII/Financial tags, and upstream/downstream lineage — explicitly marking anything DataHub doesn't provide as unavailable rather than filling gaps with guesses.
  3. Grounds Gemini 2.5 Flash in that exact metadata context and asks for structured JSON output: a dbt staging model, a dbt transformation model, a schema.yml, an Airflow DAG, data-quality tests, and documentation — with every governance decision (e.g. how PII is or isn't handled) explained explicitly, not silently applied.
  4. Validates the output against the real DataHub schema. If the model references a column DataHub never reported, validation fails and the pipeline stops — no PR gets created.
  5. Opens a real GitHub pull request — a real branch, real commits, real PR — only once validation passes.

The UI surfaces every step of this honestly: it always shows whether DataHub context came from a live MCP connection or from clearly-labeled demo fixtures, it never fakes a "success" state, and it includes a direct side-by-side of what a generic LLM produces versus what DataPilot produces with real organizational context.

How we built it

We treated this as a strict "working MVP over architectural complexity" build. The stack:

Backend: FastAPI, with a clean adapter boundary (DataHubClient) so the rest of the app never knows whether it's talking to the real MCP server or to demo fixtures — same interface, swappable implementation, controlled by one DATAHUB_MODE env var.

Agent: a plain sequential async workflow (understand → fetch metadata → generate → validate → PR) rather than a heavyweight graph framework — easier to debug in a time-boxed build, and there's no branching complex enough to need one.

LLM: Gemini 2.5 Flash via the official google-genai SDK, with structured JSON output mode and a bounded retry-and-repair loop (max 2 retries) for malformed responses.

Validation:a real hallucination detector that cross-references every column the model claims to have used against DataHub's actual schema, and a governance layer that surfaces PII/Financial handling decisions instead of blindly hashing every identifier.

GitHub integration: the real REST API — branch creation, file commits, and PR opening, gated so it's structurally impossible to open a PR when validation has failed.

Frontend: React + Vite, deliberately restrained — a request box, a metadata context panel, an agent activity feed (high-level events only, no exposed chain-of-thought), a tabbed code viewer, a validation panel, and the GitHub result.

We also wrote a seed script using DataHub's Python emitter API so a real DataHub instance can be populated with the exact same four demo datasets (raw.customers, raw.payments, raw.customer_transactions, analytics.customer_revenue) that back our fixture mode — so the live-MCP demo and the fixture demo are identical in substance, just different in data source.

Challenges we ran into

Making "never invent columns" actually enforceable, not just a prompt instruction.LLM system prompts are suggestions, not guarantees. We had to build a real validation layer that parses generated SQL, extracts referenced identifiers, and diffs them against DataHub's reported schema — and wire it so a failed diff structurally blocks PR creation rather than just logging a warning.

DataHub's MCP transport isn't a simple REST call.The official mcp-server-datahub runs over MCP's stdio transport via uvx, not plain HTTP — which meant building a real MCP client session (spawn subprocess, JSON-RPC handshake, dynamic tool discovery) rather than a quick requests.get(). We also learned MCP tool names are discovered per-server-version, not hardcoded, so we built the client to fail loudly and point at session.list_tools() rather than silently guessing wrong endpoint names.

Resisting the urge to fake anything under time pressure. It would have been trivially easy to hardcode a "success" PR URL or a canned "metadata retrieved" banner to make a demo look further along than it was. We held a hard line: if Gemini isn't configured, the app says so; if DataHub data came from fixtures, the UI says "DEMO" in the mode banner, not "live"; if validation fails, the PR button simply doesn't produce a PR. It's less flashy under a broken config, but it's honest, and honesty is the entire pitch of the product.

Keeping the demo fast without a fragile multi-container stack. DataHub's full stack (GMS + Kafka + Elasticsearch + MySQL) is too heavy to bundle reliably into a single docker compose up alongside our own app for judging. We split it: DataPilot runs in its own lightweight compose stack, and DataHub runs standalone via its own quickstart — with a demo-fixture fallback mode so the whole thing is judgeable even without a live DataHub instance running.

Accomplishments that we're proud of

A validation engine that actually catches hallucinated columns in generated SQL — not a prompt promise, a tested mechanism. An architecture where swapping demo data for a live DataHub instance requires changing exactly one environment variable, because the abstraction boundary was designed correctly from the start. A GitHub integration that's real, not simulated — genuine branches, commits, and pull requests, hard-gated behind passing validation. A product that's honest about its own limitations in its own UI — showing "DEMO" mode plainly, refusing to claim write-back succeeded when it didn't, and never dressing up a partial result as a complete one.

What we learned

The interesting engineering problem in "AI + your metadata" isn't the LLM call — it's the trust boundary around it. Anyone can prompt a model to say "I'll only use real columns." The actual work is in making that promise mechanically true: structured output you can validate, a schema diff you can run in milliseconds, and a gate that makes it impossible for a hallucinated pipeline to reach a pull request. We also relearned a boring but important lesson under time pressure: every shortcut that fakes a success state costs you the one thing a "trustworthy AI agent" product can't survive without — actually being trustworthy.

What's next for DataPilot AI: The Metadata-Aware AI Data Engineer

Write-back to DataHub— closing the loop so generated pipeline docs and lineage are pushed back into DataHub automatically, not just read from it. A real SQL AST parser(e.g. sqlglot) in place of our current identifier-extraction validator, for airtight column-reference checking against complex joins and CTEs. Multi-dataset pipelines — requests that span several source tables with independently retrieved metadata and merged lineage, not just a single source-to-target flow. Continuous validation— re-running DataPilot's validator against schema drift, so a PR that was valid on creation gets flagged automatically if the upstream schema changes before merge. Deeper governance automation — beyond flagging PII/Financial tags, generating the actual masking/encryption transformation logic when policy requires it, with the decision still explained in plain language rather than silently applied.

Built With

Share this project:

Updates