# pdml-agent

An MCP server and tool-calling agent over the [property-driven-ml](https://github.com/property-driven-ml/property-driven-ml) experiment pipeline, with a human-in-the-loop gate on anything that consumes compute and a structured trace of every call.

Property-driven ML trains classifiers against formal-logic constraints, so a run is defined by a constraint, a dataset, a differentiable logic and a seed, and produces per-epoch metrics for both predictive performance and constraint security. That makes it a genuinely tool-shaped domain rather than a demo one: experiments can be listed, configs recovered, results read, runs compared, and new runs planned, approved and executed.

Status: **complete as scoped.** Server, agent, gate, tracing. Real execution demonstrated on CPU.

## Architecture

    ┌──────────────────────────────────────────────────────────────┐
    │  agent.py                       (Anthropic SDK tool runner)  │
    │                                                              │
    │   claude-opus-5 ──► pending tool_use ──► ToolLedger.wrap     │
    │        ▲                                   │  memoise (RO)  │
    │        │                                   │  gate (compute)│
    │        │ tool_result                       │  trace (JSONL) │
    │        └───────────────────────────────────┘        │        │
    └─────────────────────────────┬───────────────────────┼────────┘
                       MCP over stdio                     ▼
    ┌─────────────────────────────┴────────────────┐   traces/*.jsonl
    │  server.py           (mcp MCPServer, thin)   │
    │   list_experiments  get_experiment_config    │
    │   get_results       compare_runs             │
    │   search_logic_definitions                   │
    │   run_experiment ──► PDML_ALLOW_EXECUTE=1 ?  │
    └────┬───────────┬──────────────┬──────────────┘
         ▼           ▼              ▼
    experiments.py  logic_defs.py  runner.py ──► subprocess: main.py
    (read CSVs)     (parse source)  (plan/execute)   in property-driven-ml

`agent.py` knows nothing about the domain. It connects to the server over stdio like any other MCP client and works only with the tools the server exposes. The domain modules carry no MCP dependency and are testable by import. `server.py` only registers tools and delegates.

## Layout

    pdml_agent/
      experiments.py   reading and comparing runs
      logic_defs.py    searching the logic implementations
      runner.py        validating, planning and executing runs
      server.py        the MCP layer, deliberately thin
      agent.py         the agent: runner, gate, memoisation, tracing
    scripts/
      make_fixtures.py generate sample runs
      smoke_test.py    start the server, exercise every tool, check refusals
      demo.py          run the agent on five tasks
    fixtures/results/  sample runs, so nothing needs a GPU to demo
    demo_output/       what the agent said and did, one JSON per task
    traces/            one JSONL per run, every turn and every call

## Tools

| Tool | Returns |
|---|---|
| `list_experiments` | runs, filterable by constraint, dataset or logic |
| `get_experiment_config` | the config a run actually trained with |
| `get_results` | metrics for one epoch, defaulting to the last |
| `compare_runs` | config and metric diff between two runs |
| `search_logic_definitions` | logic classes, their operators and docstrings |
| `run_experiment` | with `dry_run=true`, a validated command plan; with `dry_run=false`, execution behind two gates |

## The gate

`run_experiment` is the only tool that consumes compute, and two independent things stand in front of it.

**The server will not execute unless it was started with `PDML_ALLOW_EXECUTE=1`.** That is a decision made by whoever runs the server, and no request can change it. Without it, `dry_run=false` returns `status: refused` with the plan attached, and it is not an error.

**The agent will not send an execution request without an operator approving the exact call.** The approval prompt shows the tool name and the full arguments as JSON, not a summary. A decline returns a normal result reading `declined_by_operator`, and the model is instructed to report that and stop rather than retry.

Either layer alone would stop an unwanted run. Both together mean neither has to be perfect. The policy deciding what needs approval is one function, `needs_approval`, small enough to read in a glance.

## The trace

Every run appends to `traces/<timestamp>-<question>.jsonl`. One line per event, never rewritten.

`turn` records carry the step number, the model's stop reason, its text and thinking summary, the calls it is about to make, and that turn's token usage. `tool_call` records carry the tool, its arguments, whether the call succeeded, came from cache, or was gated, its latency, a summary of the result, and **the model's own stated reason**, taken from the sentence it wrote alongside the call. `gate` records carry the decision. `run_start` and `run_end` bracket it with totals.

The system prompt asks the model to state in one sentence why it is making each call, and it does. From the decline-path trace:

    turn 1  "I'll start by finding the existing YG runs to confirm identifiers."
    turn 2  "No results with those filters; let me broaden."
    turn 3  "The constraint is named `standard-robustness`. Let me get the seed-0 run's config and results."
    turn 4  "Now the dry-run plan for the requested run (matching epsilon 0.3 from the seed-0 baseline)."
    turn 5  "Plan validated. Now executing it."          ← gate: declined
    turn 6  "The training run was not executed: the operator declined ..."

That trace also caught a defect in this repo's own tools. Turn 1 got an empty result because `list_experiments` filtered on the results folder name while `run_experiment` took the class name, two vocabularies for one concept. The model recovered on its own, at the cost of a turn, and its turn-3 reason says exactly what it worked out. `list_experiments` now accepts either spelling.

## What the demos showed

Five tasks, none answerable in one call. Full transcripts in `demo_output/`, full traces in `traces/`.

**A. Best logic within an accuracy budget.** Three turns. Listed the runs, fetched all four results in one parallel turn, answered YG at 0.9981 security for 0.76 points of accuracy, and said nothing was executed.

**B. Plan a variant of an existing run.** Four turns. Fetched config, comparison and logic definition in one parallel turn, called `run_experiment` with `dry_run=true`, reported the plan and the exact command, and because a matching run existed, compared them.

**C. Compare against a run that does not exist.** Three turns. Listed first rather than guessing, confirmed STL is a real logic that simply has no run, and said so.

**D. Train, operator declines.** Six turns. Planned with a dry run first, as the tool description asks, then requested execution. The approver declined. The model reported that it was not executed and did not retry, gave the plan, and answered with what existed.

**E. Train, operator approves.** Six turns, and a real training run. Same plan-then-execute sequence; the approver accepted; the server, started with execution enabled, ran `main.py` for one epoch on CPU in 28.6 seconds and wrote `fixtures/results/standard-robustness/mnist/1/YG.csv`. The agent then called `get_results` and `compare_runs` on the new run and reported final `Test-P-Metric` 0.9160 and `Test-C-Sec-self` 0.5482. **Both match the CSV.** Unprompted, it listed the confounds against the seed-0 comparison (one epoch versus ten, delay, a deliberately weakened attack budget) and observed from the epoch-0 row that constraint security is trivially 1.0 on an untrained model and only means something alongside a converged accuracy. That is a correct reading of the metric.

That seed-1 CSV is a real run and is kept beside the synthetic fixtures on purpose. Its first line is the argv it was trained with, like every other run.

## Two things worth knowing about the data

**Epoch 0 is a pre-training evaluation.** A run configured with `--epochs 10` writes eleven rows numbered 0 to 10. Row count and final epoch are reported separately, because calling the row count "epochs" overstates training by one.

**The training script writes `-1` for metrics it did not evaluate.** `get_results` normalises those to null, so a sentinel cannot be read as a measurement. A baseline run has no constraint metrics at all, and it should say so rather than reporting minus one.

## Limits, stated so they are not overclaimed

The model never hit an `is_error` tool result live across five tasks, because it followed the instruction to list before trusting an identifier. The error path is tested at the protocol level in `smoke_test.py` and at the wrapper level, but live recovery from a mid-task tool error was not demonstrated.

Memoisation never fired live. The model did not repeat an identical call in any run. It is unit-tested and idle in every trace.

Prompt caching is not configured. `cache_read_input_tokens` is zero in every trace, and the input token counts (11k to 46k per task) are mostly re-sent context. Cache breakpoints on the tool definitions and system prompt would cut that substantially and are the obvious next improvement.

Executing a run required a checkout whose `main.py` parses. On upstream `main` it does not: `--epsilon` and `--delta` are each defined twice and argparse rejects the duplicate before any argument is read, so `python main.py --help` fails. That is fixed on the `fix/duplicate-argparse-flags` branch of the fork, with a regression test, and the demo pointed `PDML_REPO_DIR` at that checkout.

## Try it

    uv sync
    uv run python scripts/make_fixtures.py
    uv run python scripts/smoke_test.py

The smoke test starts the server over stdio, enumerates the tools, calls every one, checks that execution without `PDML_ALLOW_EXECUTE` is refused, and checks that an unknown experiment id errors rather than silently succeeding. It costs nothing.

To ask the agent something, with `ANTHROPIC_API_KEY` set:

    uv run python -m pdml_agent.agent "Which mnist run has the best constraint security?"
    uv run python scripts/demo.py A B C D

To let it actually train, point it at a checkout of property-driven-ml whose `main.py` parses and at an interpreter with torch, then pass the flag that enables execution:

    export PDML_REPO_DIR=~/property-driven-ml
    export PDML_PYTHON=~/property-driven-ml/.venv/bin/python
    uv run python -m pdml_agent.agent --allow-execute "Train a one-epoch YG run on mnist at seed 2 ..."
    uv run python scripts/demo.py E

You will be shown the exact call and asked to approve it.

Environment variables the server reads: `PDML_RESULTS_DIR` (where runs live, default `fixtures/results`), `PDML_REPO_DIR` (the property-driven-ml checkout), `PDML_PYTHON` (interpreter for `main.py`, else the repo's `.venv`), `PDML_ALLOW_EXECUTE` (`1` to permit execution), `PDML_EXECUTE_TIMEOUT` (seconds, default 3600).