Skip to main content
Glama
kishanraj41

Polygraph MCP Server

by kishanraj41
README.md
# Polygraph

![Polygraph — a lie detector for data catalogs](docs/assets/hero.png)

**A lie detector for data catalogs.**

Your catalog says the fraud model reads `raw_claims` and `legacy_claims_archive`.
Polygraph runs the pipeline, watches what it actually touches, and reports that
`legacy_claims_archive` has not been read since the refactor — while
`fee_schedule`, which nobody declared, is merged into the training set on every
run. Then it writes those findings back into DataHub as tags, so the next person
to open the catalog sees them.

Catalog says X. Runtime proves Y. Polygraph reconciles them inside DataHub's own UI.

**And it is built for agents.** Agents read catalogs over MCP now, and an agent
asking DataHub inherits the catalog's mistakes at machine speed. Polygraph ships
its own MCP server to run beside **DataHub's official `mcp-server-datahub`**:
one supplies what the catalog *claims* — testimony — the other what the runtime
*proved* — evidence — and the agent is required to keep the two apart.

```json
{
  "mcpServers": {
    "polygraph": {
      "command": "python",
      "args": ["-m", "polygraph.mcp_server"],
      "env": { "PYTHONPATH": "<repo>/src", "POLYGRAPH_GMS": "http://localhost:8080" }
    },
    "datahub": {
      "command": "python",
      "args": ["-m", "mcp_server_datahub", "--transport", "stdio"],
      "env": { "DATAHUB_GMS_URL": "http://localhost:8080" }
    }
  }
}
```

Use full interpreter paths in a real config — the tested Windows version is
[`docs/claude_desktop_config.example.json`](docs/claude_desktop_config.example.json).
A recorded run of the agent working across both servers is in
[`examples/agent_transcript.md`](examples/agent_transcript.md).

Built for [Build with DataHub: The Agent Hackathon](https://datahub.devpost.com/).

---

## The problem

Lineage in a data catalog is **testimony**. Someone wrote it down — by hand, or
via an ingestion connector that parsed SQL, or from a DAG definition. Then the
code changed and the testimony did not.

Nobody notices, because a catalog has no way to be wrong out loud. A stale edge
looks exactly like a correct one. A missing edge looks like nothing at all. Data
scientists make decisions on lineage that has quietly drifted from reality, and
the first sign of trouble is a model that stopped working for reasons no one can
trace.

Polygraph closes that loop by making the runtime testify.

---

## Verdict semantics

| Verdict | Means | Tag written |
| --- | --- | --- |
| `VERIFIED` | Declared, and runtime capture proves data flowed along it | `polygraph:verified` |
| `PHANTOM` | Declared, but nothing flowed along it in the captured run | `polygraph:phantom` |
| `UNDECLARED` | Runtime proves the edge exists; the catalog never mentioned it | `polygraph:undeclared-source` |

Read those carefully, because the asymmetry is real and Polygraph does not paper
over it:

- `VERIFIED` is evidence **from the run that was captured**. It is not a proof
  about every run.
- `PHANTOM` means nothing flowed **in this run**. A genuinely conditional edge —
  a branch not taken — will look phantom. The report always names the run it is
  based on so a human can make that call.
- An observed node with no entry in `urn_map.yaml` is reported as **unmapped**,
  never guessed at. Polygraph does no fuzzy matching.

---

## Verify this without installing anything

The quickstart needs Docker and an ~8 GB image pull. None of that is required
to check the claims — every artifact in this README ships in
[`examples/`](examples/), and checking takes about a minute:

- **The three verdicts** —
  [`reconciliation_report.md`](examples/reconciliation_report.md) is the
  reconciliation table below, exactly as the run produced it.
- **The incident report and its digest** — hash the shipped file; it matches
  the digest stored on the DataHub document:

  ```bash
  sha256sum examples/incident_report.md
  # acbedff47da6255e6b69877f722e52c2421f711e560d8517919e04bfe12ee5d3
  ```

- **The agent, actually running** —
  [`agent_transcript.md`](examples/agent_transcript.md) is a recorded
  `ask --llm` run over both MCP servers. Gate 11 asserts it called real tools
  on both and cited the numbers those tools returned.
- **What it looks like in DataHub** —
  [`ui_undeclared_source.png`](examples/ui_undeclared_source.png),
  [`ui_undeclared_no_lineage.png`](examples/ui_undeclared_no_lineage.png),
  [`ui_lineage.png`](examples/ui_lineage.png),
  [`ui_integrity_score.png`](examples/ui_integrity_score.png),
  [`ui_incident_document.png`](examples/ui_incident_document.png) — the tags,
  the empty lineage tab, the structured property, and the incident document,
  on screen.

Everything in `examples/` except the agent transcript reproduces byte-for-byte
from a clean clone. The transcript says so itself: model output varies between
runs.

---

## Quickstart

Prerequisites: Docker Desktop with ~8 GB allocated, Python 3.11+, ~13 GB free
inside the Docker VM.

```bash
git clone https://github.com/kishanraj41/datahub-polygraph
cd datahub-polygraph
python -m venv .venv && . .venv/bin/activate    # Windows: .venv\Scripts\activate
pip install -r requirements.txt

# 1. Stand up DataHub (first run pulls ~8 GB, 10-25 min)
datahub docker quickstart
datahub init --username datahub --password datahub

# 2. Seed the catalog with deliberately imperfect lineage
python demo/seed_catalog.py

# 3. Run the pipeline under AutoLineage capture
python demo/pipeline.py --mode healthy

# 4. Reduce the capture to a dataset-level graph
python -m polygraph.cli observe \
    --trace runs/healthy/trace.json \
    --out   runs/healthy/observed_graph.json --root .

# 5. Reconcile declared against observed
python -m polygraph.cli reconcile --observed runs/healthy/observed_graph.json

# 5b. Read catalog context (owners, descriptions, search) through DataHub's
#     own MCP Server. Optional; needs DataHub running.
python -m polygraph.cli catalog --search "/q fee+schedule"

# 6. Write the verdicts back into DataHub
python -m polygraph.cli writeback \
    --report   examples/reconciliation_report.json \
    --document examples/reconciliation_report.md
```

Open <http://localhost:9002> (login `datahub` / `datahub`) and look at
`polygraph.demo.fee_schedule`. It is tagged `polygraph:undeclared-source`.

On Windows, `scripts/run_gate1.ps1` then `scripts/run_gate2.ps1` do all of the
above with preflight checks.

### The incident path

```bash
python demo/pipeline.py --mode buggy     # one changed line collapses F1
python -m polygraph.cli observe --trace runs/buggy/trace.json \
    --out runs/buggy/observed_graph.json --root . --mode buggy
python -m polygraph.cli incident
```

---

## Real output

Everything below is copied from actual runs, reproduced on two machines
(Linux/Python 3.11 and Windows 11/Python 3.12) with the pinned dependencies in
`requirements.txt`. Full artifacts are in [`examples/`](examples/).

**Reconciliation** — one of each verdict against the seeded catalog:

| Verdict | Upstream | Operations observed |
| --- | --- | --- |
| `VERIFIED` | `polygraph.demo.raw_claims` | filter → concat → merge → drop → select → LogisticRegression.fit |
| `PHANTOM` | `polygraph.demo.legacy_claims_archive` | — |
| `UNDECLARED` | `polygraph.demo.fee_schedule` | filter → merge → drop → select → LogisticRegression.fit |

**Incident** — one changed line (`quantile(0.999)` → `quantile(0.05)`):

| | baseline | degraded |
| --- | ---: | ---: |
| F1 | 0.8282 | 0.0000 |
| rows after filter | 5994 | 300 |

AutoLineage — the author's own capture library (see
[Prior work](#prior-work)) — localises the collapse to the **`filter`**
operation with an impact score of 1.0, against a next-ranked deviation three
orders of magnitude lower. The incident document names the owning team
(`urn:li:corpGroup:ml-platform-team`) resolved live from DataHub ownership.

The document's sha-256 is stored on the DataHub document as `polygraph_sha256`.
The report is byte-reproducible: rerunning the incident path on the same code
produces the identical file and therefore the identical digest. You can check
both — hash the shipped file, then regenerate it and hash again:

```bash
sha256sum examples/incident_report.md
# acbedff47da6255e6b69877f722e52c2421f711e560d8517919e04bfe12ee5d3
```

---

## Ask an agent instead

`mcp-server-datahub` lets an agent read what the catalog **claims**. Polygraph
ships its own MCP server so the same agent can read what the runtime **proved**,
and the gap between them.

```bash
python -m polygraph.mcp_server        # stdio
```

Register it alongside DataHub's own server — the config block at the top of
this README registers both, and the tested Windows version is
[`docs/claude_desktop_config.example.json`](docs/claude_desktop_config.example.json).

| Tool | Answers |
| --- | --- |
| `can_i_trust(asset_urn)` | Did this asset's declared lineage survive a real run? |
| `get_integrity_score(job_urn)` | Score, precision, recall, and which way the catalog is wrong |
| `list_undeclared_sources()` | What does the pipeline read that nobody declared? |
| `list_phantom_edges()` | Which declared edges carried no data? |
| `get_incident_report(urn)` | The hash-verified incident, root operation and owner |
| `explain_verdict_semantics()` | What the verdicts do **not** establish |

Two design choices worth naming, because they are what stop an agent from
overstating the findings:

**Absent evidence is never silence.** Every tool returns `evidence_available`.
Asked about an asset Polygraph has not examined, `can_i_trust` says so
explicitly — *"there is no evidence either way. Do not treat this as a clean
bill of health."* An empty result that reads as "nothing wrong" is the failure
mode this project exists to complain about, so the server refuses to produce
one.

**Every tool returns evidence, not just a verdict.** An agent that relays the
verdict is correct; one that reads the operation paths can disagree with it.
`explain_verdict_semantics` exists so an agent can look up what a verdict does
*not* establish before repeating it to a person.

### Or ask from the command line

```bash
polygraph ask "what undeclared sources does the pipeline read?"
polygraph ask "can I trust fee_schedule?"
polygraph ask "why did f1 drop?"
```

Two backends over the **same six tool functions** — there is one implementation
of "can I trust this asset", not two:

**Deterministic (default).** A keyword router. It needs no API key and produces
identical output for identical input, which is why every claim in this README
reproduces from a bare clone. It is not an agent and does not describe itself as
one. Asked something it cannot classify, it declines and lists what it can
answer rather than running the nearest-matching tool — silently answering a
different question than the one asked is worse than declining. Exit code 3 when
it does not understand, so it is scriptable.

**LLM (`--llm`).** A real tool-use loop over **eight** tools — Polygraph's six,
plus `datahub_get_entities` and `datahub_search`, which proxy to DataHub's own
MCP Server. The model picks tools and writes the answer from what they return.
Needs `ANTHROPIC_API_KEY` and `pip install anthropic`. Gated behind a flag on
purpose: a judge should never need credentials to verify a documented result.
Without a key it says so plainly rather than degrading into something else.

The system prompt requires the model to say **"Polygraph observed"** for
anything from a Polygraph tool and **"the catalog says"** for anything from a
`datahub_*` tool. The gap between evidence and testimony is the whole subject;
an agent that blurs it in its answer would defeat the point of asking.

A real run is recorded in
[`examples/agent_transcript.md`](examples/agent_transcript.md) — produced by
`scripts/run_gate11.ps1`, which asserts the agent actually called tools, reached
both servers, and cited the numbers those tools returned. It is the one file in
`examples/` that is **not** reproducible, and it says so at the top: language
model output varies between runs. Everything else there reproduces byte for
byte.

```bash
export POLYGRAPH_ASK_MODEL=claude-opus-5      # or --model; the default rots
polygraph ask --llm "Polygraph says the pipeline reads an undeclared source. \
  Which asset is it, is it registered in the catalog, and who owns it?"
```

The model ID is deliberately overridable. This shipped with a hardcoded ID that
had been **retired for two months**, so `--llm` failed every time — and nothing
noticed, because nothing ran it. That is what Gate 11 is for.

### Talking to DataHub through DataHub's own MCP Server

Polygraph reaches DataHub two ways: the `acryl-datahub` SDK, and
**`mcp-server-datahub`** — DataHub's own MCP Server, launched as a stdio
subprocess exactly the way an agent client launches it.

That is deliberate rather than decorative. Polygraph's argument is about what a
catalog tells the people and agents who ask it, so it should be checking the
answer DataHub actually gives an agent.

Three MCP tools are used:

| Tool | What Polygraph asks it | Where |
|---|---|---|
| `get_entities` | who owns this asset, how is it described | `polygraph catalog`, `ask --llm` |
| `search` | is this undeclared source registered at all | `polygraph catalog --search`, `ask --llm` |
| `get_lineage` | what does the catalog say feeds this job | `reconcile --declared-via mcp` |

`get_lineage` is **not** the default path. The reason is a story worth keeping:
on the stack this was built against it returned a 500 from GMS, and the first
diagnosis written here — a search-dialect misconfiguration — was **wrong**. The
OpenSearch container had simply died, and GMS could not resolve the hostname
`search`. A DataHub stack in that state answers `/config`, reports healthy,
serves every entity read, and fails only the queries that need search. See
[`docs/DATAHUB_MCP.md`](docs/DATAHUB_MCP.md) for the full account, including what
would have caught it sooner.

The default stays on the SDK for the reason that outage demonstrated: it reads
the `dataJobInputOutput` aspect from MySQL and kept working the whole time, while
every search-backed path was down. A default should be the one that still answers
when something is broken.

When the MCP lineage path *does* work, `scripts/run_gate10.ps1` reconciles twice,
once per path, and fails if the per-edge verdicts differ. Matching totals with
differing edges fails too — a summary-only check would let that through.

Every MCP response is parsed by **walking** the payload rather than indexing a
fixed path. A hard-coded path would break silently on a server upgrade: an empty
upstream set would make Polygraph report a catalog full of phantom edges, and a
missing `ownership` node would make it report an unowned asset. Confidently
wrong, in the exact way this project exists to catch. An empty lineage result
therefore raises instead of producing verdicts.

**Not-found is reported in-band, and reading it is not optional.** `get_entities`
runs an existence check and, for a URN it cannot find, returns
`{"error": "Entity <urn> not found", "urn": ...}` — an entry that still carries a
`urn`. Index the response by URN and treat presence as existence, and a
fabricated asset comes back as a real one. Polygraph honours the server's verdict
and reports the URN as `found: false` with the reason attached, because "no owner
recorded" and "no such asset" are different answers. Gate 10a checks it against
the live catalog with a deliberately fabricated URN.

A second finding, separate from the point-in-time bug: the MCP Server **cannot**
report a data job's declared inputs at all. `inputOutput` / `inputDatasets` /
`outputDatasets` appear in none of its GraphQL documents, though GMS exposes
`DataJob.inputOutput` — verified by `scripts/probe_gms.py`. That is a one-field
upstream fix, not a DataHub limitation.

---

## Architecture

```mermaid
flowchart LR
    subgraph Catalog["DataHub"]
        DJ["dataJob<br/>train_fraud_model"]
        DS["datasets"]
        KB["knowledge base"]
    end

    subgraph Runtime["Your pipeline, unmodified"]
        PIPE["pandas + scikit-learn"]
        AL["AutoLineage<br/>239 import-time hooks"]
    end

    PIPE -->|"captured by"| AL
    AL -->|"trace.json"| OBS["observed.py<br/>operation graph → dataset graph"]
    DJ -->|"dataJobInputOutput"| DEC["declared.py"]

    OBS --> REC["reconcile.py<br/>VERIFIED / PHANTOM / UNDECLARED"]
    DEC --> REC
    REC --> WB["writeback.py"]
    WB -->|"tags"| DS
    WB -->|"documents"| KB

    REC --> INC["incident.py<br/>metric delta + root-cause ranking"]
    INC --> KB
```

The interesting problem is the middle box. AutoLineage records lineage at
*operation* granularity — every pandas transform, every sklearn call, linking
dataframe versions. DataHub declares lineage at *dataset* granularity. Bridging
them required three non-obvious decisions, documented in
[`src/polygraph/observed.py`](src/polygraph/observed.py):

1. **Anchors.** A node is catalog-visible if it is a file read, a file write, or
   a fitted model. Two anchors form an edge when a path connects them through no
   other anchor.
2. **Union of paths, not shortest path.** AutoLineage links
   `LogisticRegression.fit` directly back to an early hub node, so the shortest
   route from source file to model skips the filter, the merge and the split.
   Those operations really ran between the two assets. Shortest-path would have
   left the incident report with nothing to name.
3. **Self-loops carry the payload.** When a transform does not change dataframe
   identity, AutoLineage emits `parent_id == child_id`. The planted bug *is* one
   of those. Discarding them is correct for topology and catastrophic for
   diagnosis.

---

## Limitations

Stated plainly, because a tool that accuses catalogs of lying should be direct
about what it cannot do.

- **Single-run evidence.** Every verdict describes one captured run. Conditional
  branches not taken during capture are indistinguishable from dead edges.
- **Inputs only.** Polygraph reconciles edges *into* a job. It does not
  reconcile outputs. AutoLineage cannot link a numpy `predict()` output into a
  newly constructed DataFrame, so a declared `job → predictions` edge would come
  back `PHANTOM` — a tool limitation dressed up as a stale catalog edge. Scoping
  to inputs is also what the tags claim: `undeclared-source`.
- **Shape-preserving bugs are invisible.** The analyzer ranks row-count and
  column-count deviations. A unit error that scales a column by 1000 while
  preserving every shape would not appear at all.
- **Localisation is to an operation, not a line number.**
- **Info-level anomalies are excluded from the incident report.** AutoLineage
  emits timing-sensitive counters at `info` severity that vary between identical
  runs. Including them made the report non-reproducible — same code, same seed,
  different sha-256. They are filtered so the published digest means something;
  they do not affect localisation.
- **Synthetic demo data.** `demo/pipeline.py` generates its own 6,000-row
  dataset with a fixed seed so the demo reproduces from a clean clone with no
  downloads. The AutoLineage paper's headline case uses the real Kaggle
  credit-card fraud dataset; that data is 150 MB and not redistributable, so it
  is not part of this repo.
- **`pathlib.Path` breaks capture.** AutoLineage's IO hooks test
  `isinstance(path, str)`, so `pd.read_csv(Path(...))` records *nothing* — no
  file lineage at all, silently. `demo/pipeline.py` passes `str(...)`
  explicitly. This is an upstream bug, not a design choice.
- **Everything search-backed depends on a healthy OpenSearch.** That includes
  `reconcile --declared-via mcp`, `polygraph catalog --search`, the
  `datahub_search` agent tool, and DataHub's own UI search and Lineage tab. When
  OpenSearch is down the rest of DataHub keeps answering, so this is easy to
  mistake for a Polygraph bug. `scripts/stack_status.ps1` settles it in seconds.
- **DataHub's MCP Server cannot report declared job inputs.** Its GraphQL
  documents never request `DataJob.inputOutput`. So even with the point-in-time
  bug fixed, the only MCP route to declared lineage is `get_lineage`, which
  reports the *rendered* graph rather than the asserted aspect.
- **Catalog context is unverified.** `polygraph catalog` and the `datahub_*`
  agent tools report owners and descriptions as the catalog states them.
  Polygraph verifies lineage and nothing else — an owner field can be as stale
  as an edge, and Polygraph will not tell you.
- **Existence comes from DataHub, not from inference.** Polygraph reports an
  asset as missing only when the MCP Server's own existence check says so. It
  does not guess from how sparse a response looks — an entity can be registered
  and entirely undocumented, and calling that nonexistent would be as wrong as
  inventing one.
- **Python 3.12 is untested by DataHub.** The CLI warns. It worked throughout
  this build, but 3.11 is the supported version.
- **One pipeline, one catalog shape.** The URN mapping is explicit YAML. Using
  Polygraph on your own pipeline means writing your own `urn_map.yaml`.

---

## Prior work

Polygraph is built on the author's own earlier work, disclosed here rather than
left for a reader to discover:

1. **AutoLineage: Zero-Code Data Lineage for Python ML Pipelines** — the runtime
   capture library Polygraph depends on, and the source of the planted-bug
   evaluation this demo's fraud pipeline is drawn from.
   [SSRN 6683825](https://papers.ssrn.com/sol3/papers.cfm?abstract_id=6683825)
2. **RudriQ** — deviation-weighted root-cause analysis and deterministic audit
   reporting, which the incident path here follows.
   [SSRN 7222278](https://papers.ssrn.com/sol3/papers.cfm?abstract_id=7222278)

`autolineage` is MIT-licensed and installed from PyPI as a pinned dependency.
No code from either project was copied into this repository.

What is new here, and what is not: the runtime capture and the anomaly ranking
are prior work. The reconciliation of *declared* against *observed* lineage, the
three verdicts and their semantics, the integrity score, the write-back into
DataHub, and both MCP integrations were written for this project.

---

## License

Apache License 2.0. See [LICENSE](LICENSE).