Part of the PolicyTrace series

Document AI

My tests passed through a model swap. That was the problem.

PolicyTrace's tests passed before and after I changed its model, because none of them call a model. What's tested, what isn't, and the smallest evaluation I'd trust.

Teja Sagiraju · 4 min read

When Groq retired the Llama models PolicyTrace was built on, I switched the extraction to openai/gpt-oss-120b and the classifier to openai/gpt-oss-20b. The arbiter tests passed before the change and passed after it, in about 0.2 seconds. They would also have passed if the new model had got every premium wrong, because none of them ever calls a model.

So how did I decide the new model was fine? I ran the demo pack and read the output. That's an honest answer, and it isn't good enough. This post is about what PolicyTrace tests today, what it doesn't, and the smallest evaluation I think it needs.

What is tested

Diagram of the PolicyTrace pipeline stages, Docling, PIIMasker, classify, extract, PolicyArbiter, provenance and review UI, where only PolicyArbiter has automated tests; manifest.json with three expected values points at the extract stage
At the time of the model swap, one stage out of seven had automated tests.

At the time of the swap, every test was in tests/test_arbiter.py, and they're the right kind of tests for that code. The arbiter is pure Python, so the tests build two small UKMotorGoldenRecord objects by hand, merge them and check the result:

  • the Schedule's vehicle, premium and risk data are kept;
  • the Certificate wins on class of use;
  • a policy number that differs produces a ConflictEntry with both values and the winner;
  • "ALICE J SMITH" on one document and "ALICE SMITH" on the other are matched as one driver;
  • two values that differ only by whitespace are not a conflict;
  • an explicit False for main driver on the Schedule is kept, and a missing one is filled from the Certificate;
  • for every conflict, the merged record holds the winner's value.

The last three points, four tests in all, came out of one live run, written up in Two documents, one policy, four disagreements. That's the pattern I'd like to keep: when a run shows a problem, the fix lands with a test that fails on the old code.

What isn't

Everything before and after the arbiter was checked by eye. Nothing tested that Presidio masks the policyholder's name, that the classifier labels the four demo PDFs correctly, that the extraction gets the right excesses, or that a highlight lands on the value. On the latest run, a handful of driver fields were highlighted on a heading, a label or another driver's row. The driver merge also paired drivers with the wrong person, because masked names are numbered separately in each document. Both are exactly the kind of bug an automated check would flag on its first run.

The closest thing to an evaluation set is sample_data/policytrace_demo_pack/manifest.json. Next to the file list it has three expected values: the policy number NBM-DEMO-0427, the registration ZX24 DEM and the insurer, Northbridge Mutual Motor Insurance Ltd. That's a good place to start, and it's still only three fields out of about sixty.

The smallest evaluation I'd trust

I don't want an evaluation framework. I want one script that runs the real pipeline on the demo pack and compares the result to a file I wrote by hand. It would need four things.

An expected record. I generated the demo pack myself with scripts/generate_synthetic_policy_pack.py, so every true value is already in that script's POLICY dictionary. The expected Golden Record can be built from it instead of typed out.

Field-aware comparison. Policy number, registration and booleans should match exactly after trimming. Dates need a normalised comparison, because the record holds ISO dates and the PDF holds "15/04/2026". Money needs one too, so that "GBP 395.00" and 395.0 count as the same. And a missing field should be reported separately from a wrong one, because they have different fixes.

The conflicts I expect. I know which disagreements between the Schedule and the Certificate this pack should produce. The script should check that exactly those conflicts appear, with the right winner, and nothing else. The two whitespace conflicts from the last run would have failed this check straight away.

A check on the evidence. For each located field, is the highlight on the right document, and does the matched text contain the value? The latest run reports "47 of 59 fields located". I'd like it to report how many highlights actually point at the value.

A run of that script costs a few minutes of Groq's free tier, which is too slow and too rate-limited for every commit. I'd run it by hand whenever I change a prompt in config/prompts.yaml, the model names in config/settings.yaml, the schema or the matcher. Those are the changes that can break extraction without breaking a single unit test, and the model retirement was exactly that kind of change.

One pack is not a benchmark

Even with all of that, the result is one synthetic policy. It tells me whether a change made things worse on a pack I know well. It doesn't tell me how PolicyTrace would do on a real insurer's layout, a scanned certificate or a pack where the drivers are listed in a different order on each document. For that I'd need more packs, and the generator script is the cheapest way to get them: vary the drivers, the excesses and the table layout, and keep the expected record in step.

Update, later the same day. I fixed the two bugs above and added tests for both. tests/test_masking.py checks that one person gets the same masked token in every document, and tests/test_provenance.py rebuilds the demo's driver table and checks where each field lands. The suite is now 32 tests: 24 for the arbiter, 5 for provenance, 2 for masking and 1 that keeps debug output off by default. Every new one fails on the old code. The full-pipeline evaluation described above is still to do.

That's what I'd like to build next. If you want to see where it will plug in, the tests are in tests/ and the demo pack is in sample_data/.

Written by Teja Sagiraju

AI engineer with 14 years across machine learning and applied AI: prompt design, document AI, LLM workflows and the review loops that make them safe to ship. I build real systems and write up what breaks along the way.

Connect

Get the next note by email

Short, first-person notes from real runs. Free.