Add uk_python typed tools (run_economy_simulation) to fix B1 reform failures

backend/agent_tools.py

@@ -729,6 +729,19 @@ def execute_tool(
         "compute": compute,
         "generate_chart": generate_chart,
     }
+    if tool_name == "run_economy_simulation":
+        from tooling.backends import get_backend as _get_typed_backend
+        try:
+            engine = _get_typed_backend(backend_id)
+        except ValueError as exc:
+            return {"error": f"{tool_name} not supported for backend={backend_id!r}", "detail": str(exc)}
+        if not hasattr(engine, "run_economy_simulation"):
+            return {"error": f"{tool_name} not implemented for backend={backend_id!r}"}
+        try:
+            return engine.run_economy_simulation(**tool_input)
+        except Exception as exc:
+            logger.exception(f"[TOOLS] {tool_name} failed")
+            return {"error": str(exc), "type": type(exc).__name__}
     if tool_name not in tools:
         return {"error": f"Unknown tool: {tool_name}"}
     try:
@@ -749,7 +762,7 @@ def execute_tool(
 
 def get_tool_definitions(backend_id: str = "uk_compiled") -> List[Dict[str, Any]]:
     backend = get_backend(backend_id)
-    return [
+    defs: List[Dict[str, Any]] = [
         {
             "name": "run_python",
             "description": backend.tool_description(),
@@ -768,6 +781,28 @@ def get_tool_definitions(backend_id: str = "uk_compiled") -> List[Dict[str, Any]
             },
         },
     ]
+    if backend_id == "uk_python":
+        from tool_definitions import RUN_ECONOMY_SIMULATION_INPUT_SCHEMA
+        defs.append(
+            {
+                "name": "run_economy_simulation",
+                "description": (
+                    "Run a UK economy-wide microsimulation comparing baseline "
+                    "current law to a parametric reform via policyengine_uk. "
+                    "Preferred over run_python for any society-wide reform "
+                    "analysis: removes the need to author Microsimulation+reform "
+                    "code, pins methodology, and returns identical numbers to "
+                    "PE-API's /uk/economy endpoint. Reform keys are programmes "
+                    "(income_tax, national_insurance, child_benefit, ...) with "
+                    "field/value subkeys (personal_allowance, basic_rate, "
+                    "main_rate, ...). Use run_python only for parameters this "
+                    "tool's mapping table does not yet cover, or for structural "
+                    "reforms."
+                ),
+                "input_schema": RUN_ECONOMY_SIMULATION_INPUT_SCHEMA,
+            }
+        )
+    return defs
 
 
 TOOL_DEFINITIONS = get_tool_definitions("uk_compiled")

backend/model_backends.py

@@ -216,31 +216,68 @@ def _ensure_importable(self) -> None:
         )
 
     def prompt_context(self) -> str:
-        return """CRITICAL - USE THE POLICYENGINE UK PYTHON MODEL INTERFACE:
-- The selected backend is `uk_python`, the Python `policyengine-uk` model package.
-- This is the detailed PolicyEngine Core/OpenFisca-style UK model, not the compiled Rust wrapper.
-- The Python environment preloads:
-  `policyengine_uk` as `pe`
-  `Simulation`
-  `Microsimulation`
-  `CountryTaxBenefitSystem`
-  `Scenario`
-  `capabilities`
-  `pd`, `np`, `json`, `math`
-- If installed, the higher-level `policyengine` package is also preloaded as `policyengine`.
-- Prefer writing code against `policyengine_uk` objects and formulas rather than recreating policy logic.
-
-COMMON WORKFLOWS FOR THIS BACKEND:
-- First inspect backend details:
-  `result = capabilities()`
-- Custom household/situation run:
-  `sim = Simulation(situation={...})`
-  `result = sim.calculate("household_net_income", 2025).tolist()`
+        return """TOOL ROUTING - READ THIS FIRST:
+
+Society-wide / economy-wide / population-level reform question?
+  → CALL `run_economy_simulation` ON YOUR FIRST TURN. Do not call
+    `run_python` first to "explore the API" or "check the interface".
+    Do not call `capabilities()`. Go straight to `run_economy_simulation`.
+
+Reform expressible in the programme/field shape (see list below)?
+  → `run_economy_simulation` is REQUIRED. Falling back to `run_python`
+    for an expressible reform is a bug: it produces drifting numbers
+    across runs and wastes 10-30 tool calls relearning the engine.
+
+Reform NOT in the mapped field list, OR question isn't about a reform?
+  → Then and only then, use `run_python`.
+
+`run_economy_simulation` shape:
+- args: `{year, reform, dataset}`. `year` defaults to 2025; `dataset`
+  defaults to "efrs" (Enhanced FRS, the same data PE-API uses).
+- `reform` is a programme→field→value dict, NOT dotted parameter paths.
+- Returns a structured dict: `budget.{budgetary_impact, tax_revenue_impact,
+  benefit_spending_impact}`, `decile.{average, relative}` keyed 1-10,
+  `poverty.poverty.{all, child, adult, senior}.{baseline, reform}`.
+- Numbers match PE-API's `/uk/economy` endpoint by construction (same
+  engine, same methodology).
+
+Worked examples (copy this shape):
+- Raise personal allowance to £15,000:
+  `{"reform": {"income_tax": {"personal_allowance": 15000}}}`
+- Basic rate 20% to 21%:
+  `{"reform": {"income_tax": {"basic_rate": 0.21}}}`
+- NI main rate 8% to 6%:
+  `{"reform": {"national_insurance": {"main_rate": 0.06}}}`
+- Stacked: basic+higher rates + NI cut:
+  `{"reform": {"income_tax": {"basic_rate": 0.22, "higher_rate": 0.42},
+               "national_insurance": {"main_rate": 0.06}}}`
+
+Mapped programme/field combinations (ONLY these are expressible):
+- income_tax: personal_allowance, basic_rate, higher_rate, additional_rate
+- national_insurance: main_rate, primary_threshold
+- child_benefit: eldest_amount, additional_amount
+
+If your reform uses any other field, the tool will return an error
+listing valid fields. THEN fall back to `run_python`.
+
+ABOUT THIS BACKEND:
+- Selected backend is `uk_python`, the Python `policyengine-uk` model package.
+- This is the PolicyEngine Core / OpenFisca-style UK model, not the compiled
+  Rust wrapper.
+
+PYTHON ENVIRONMENT (for `run_python` fallback only):
+- Preloaded: `policyengine_uk` as `pe`, `Simulation`, `Microsimulation`,
+  `CountryTaxBenefitSystem`, `Scenario`, `capabilities`, `pd`, `np`,
+  `json`, `math`. The higher-level `policyengine` package is preloaded
+  when installed.
+- Custom household: `sim = Simulation(situation={...})`,
+  `result = sim.calculate("household_net_income", 2025).tolist()`.
 - Microsimulation from published UK data:
   `sim = Microsimulation(dataset="hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5")`
-  `result = sim.calculate("household_net_income", 2025).head().to_list()`
-- Parameter reform:
-  pass parameter changes through `Scenario` or mutate a simulation with documented `policyengine_uk` helpers.
+- Parameter reform NOT expressible via `run_economy_simulation`:
+  pass a dotted-path `reform=` dict to `Microsimulation(...)`:
+  `reform = {"gov.hmrc.income_tax.allowances.personal_allowance.amount": {"2025-01-01.2025-12-31": 15000}}`
+  `sim = Microsimulation(dataset=..., reform=reform)`
 
 MODELLING SCOPE:
 - This backend exposes the Python `policyengine-uk` model surface. Its API, datasets, variables, and results can differ from `uk_compiled`.