feat: ground chat agent in a cached, auto-generated PolicyEngine API reference

[vahid-ahmadi]

Summary

Replaces drift-prone hardcoded API examples in the system prompt with a freshly generated, prompt-cached reference extracted from the installed library.

Commit 1 — inject cached reference (2084bf5)

• New backend/scripts/build_reference.py walks the installed policyengine_uk_compiled, dumps docstrings, capabilities() output, and the full Parameters JSON schema to backend/reference.md (~30k tokens).
• backend/routes/chatbot.py loads it at import and sends it as a second cache_control=ephemeral system block alongside SYSTEM_PROMPT.
• _select_chat_model counts the reference in the Haiku→Sonnet routing estimate.

Commit 2 — drop stale prompt blocks (0214aa4)

• Deletes COMMON WORKFLOWS, MODELLING SCOPE, and DATASETS sections from SYSTEM_PROMPT — all hardcoded examples now covered by the reference.
• Replaces them with a short paragraph pointing the agent at the attached reference.
• Every behavioral guardrail preserved (always compute with Python, reproducibility rules, analytical notes, user-facing style).

Why

Claude's training-data memory of the PolicyEngine API is stale the moment the library ships a release. The hardcoded examples in SYSTEM_PROMPT drift the same way and need a human edit to sync. Together these two changes eliminate both drift sources: the version-sensitive facts move to a file regenerated from the live library, and the prompt keeps only the behavioral rules. Prompt caching keeps the per-turn cost of the ~30k-token reference at ~10% of sending it fresh every time.

Measured impact

Local smoke test on Haiku 4.5:

• Fresh question in a new session → cache_creation_input_tokens ≈ 70k (reference + prompt + tools cached).
• Follow-up within 5 min → cache_read_input_tokens ≈ 70k, cache_creation_input_tokens = 0.
• Agent correctly enumerates datasets and programmes using the cached capabilities() snapshot without the deleted DATASETS / MODELLING SCOPE blocks.

Refreshing the reference

Re-run after any policyengine-uk-compiled bump:

docker-compose exec backend python scripts/build_reference.py

A follow-up PR should wire this into pr-beta-deploy.yml so it runs automatically on every library upgrade.

Test plan

• docker-compose up → send one chat message → second message should show cache_read_input_tokens > 0
• Ask "What datasets are available?" → answer lists FRS, EFRS, LCFS, SPI, WAS from the reference / live capabilities() call, not from the deleted DATASETS block
• Ask "Raise the personal allowance to £15k and show the distributional effect" → agent constructs Parameters.model_validate({...}) correctly against the current schema without the hardcoded example
• Regression: quantitative answers still compute via Python, not memory

🤖 Generated with Claude Code

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
policyengine-uk-chat	Ready	Preview, Comment	May 6, 2026 10:57am

[github-actions]

Beta preview has been cleaned up because this PR was closed.

[SakshiKekre]

I like the idea here. One thought I had is that right now the app is still consuming an already-generated reference.md. If the follow-up PR wires regeneration into the build / startup path so the backend always comes up serving a freshly generated reference, then I think this concern mostly goes away.

Otherwise, it feels like we may just be recreating the same drift issue in a different form: not stale handwritten prompt examples, but a stale generated reference artifact.
Looking forward to the follow-up PR for that piece!

One separate concern: the generator currently seems to include inherited/generic docs as well as real package docs. In the generated reference I can already see built-in container docstrings and large chunks of generic Pydantic ModelMetaclass boilerplate. That increases prompt size/cost and lowers signal.
Could we tighten the generator so it prefers actual PolicyEngine API surface and skips framework/builtin boilerplate?

[vahid-ahmadi]

Thanks @SakshiKekre — both points fair, pushed 8fb456d on this branch:

Drift (build/startup regeneration). Added RUN python scripts/build_reference.py to backend/Dockerfile after the COPY / pip install steps. Every image build now regenerates reference.md against the installed policyengine-uk-compiled version, so production always ships a fresh reference. The checked-in reference.md is from here on just a dev-time artifact for reviewing generator output without a docker build.

Generator noise. Tightened build_reference.py:

• _from_package() — classes/functions whose __module__ is outside policyengine_uk_compiled.* are skipped entirely (drops stdlib / pydantic re-exports).
• _own_doc() — compares the object's docstring to pydantic.BaseModel.__doc__ and to each base class's doc; discards matches. This is what was emitting ~30 lines of ModelMetaclass boilerplate on every Pydantic model.
• Module-level constants (DATASETS, HOUSEHOLD_DEFAULTS, BENUNIT_DEFAULTS, PERSON_DEFAULTS) now emit the value repr instead of tuple.__doc__ / dict.__doc__ — the actual contents are the signal.

Will share the post-rebuild size delta once the preview redeploys. Expected shape: the reference section drops from ~3100 lines to roughly a third of that, dominated now by real signatures + the Parameters JSON schema.

[SakshiKekre]

Thanks for the follow-up. I think this is moving in the right direction.
Question: Do we need to do this for Modal as well since our deployed backend path is Modal?

Also, the committed backend/reference.md still contains the builtin dict docs and Pydantic BaseModel boilerplate, not sure if the tightened generator was run before committing.

[vahid-ahmadi]

Thanks @SakshiKekre — both fair. Pushed 2d92139:

Modal regeneration. You're right, Modal was the gap. modal_app.py only had add_local_dir("backend", …), which copies the local checked-in tree into the image — so Modal was shipping the stale committed reference.md regardless of what the Dockerfile did. Added .run_commands("cd /app/backend && python scripts/build_reference.py") after the local-dir step, so Modal now regenerates against its own pip-installed policyengine-uk-compiled. The Dockerfile already did this; both deploy paths are now symmetrical.

Stale committed reference. You diagnosed it correctly — the committed file was generated before the tightened generator landed, which is why it still has the ModelMetaclass boilerplate. Rather than re-running and re-committing (which would just reset the drift timer), I removed backend/reference.md from git entirely and added it to .gitignore. Both deploy paths now build it at image-build time, so there's no production reason to ship a static copy, and no more drifting artifact for reviewers to spot. Local dev regenerates with docker-compose exec backend python scripts/build_reference.py or by running the script directly.