Skip to content

feat: add structural reform support for AI agent#67

Merged
nikhilwoodruff merged 5 commits intomainfrom
feat/structural-reforms
Jan 8, 2026
Merged

feat: add structural reform support for AI agent#67
nikhilwoodruff merged 5 commits intomainfrom
feat/structural-reforms

Conversation

@nikhilwoodruff
Copy link
Contributor

@nikhilwoodruff nikhilwoodruff commented Jan 8, 2026

Adds ability for the AI agent to write and run structural reforms (simulation modifiers) on UK/US simulations.

Closes #66

Changes

Data model:

  • Added simulation_modifier field to Policy model for storing Python code defining custom variable formulas
  • Database migration to add the column

Modal functions:

  • Updated _get_pe_policy_* functions to return tuple of (policy, simulation_modifier)
  • Added _apply_simulation_modifier() helper to execute modifier code
  • Updated all economy simulation/comparison functions to apply modifiers

Agent upload endpoints (/agent/results/*):

  • POST /simulation - mark simulation completed/failed
  • POST /household - upload household calculation results
  • POST /decile-impact - upload decile impact data
  • POST /program-statistics - upload program statistics
  • POST /poverty - upload poverty metrics
  • POST /inequality - upload inequality metrics
  • POST /report-status - update report status
  • POST /policy-with-modifier - create policy with simulation modifier

Agent tools:

  • Added execute_python tool so the agent can test modifier code before submitting
  • Runs Python in a sandboxed environment, returns stdout/stderr/exceptions

Agent system prompt:

  • Added workflow for structural reforms with example modifier code
  • Instructions to always test code with execute_python before submitting

Example usage

The agent workflow for structural reforms:

# 1. Test the modifier code first
execute_python(code="""
from numpy import where

def modify(simulation):
    pass

print(f"modify defined: {callable(modify)}")
# Test logic
income = 15000
benefit = 1000 if income < 20000 else 0
print(f"Test: income={income} -> benefit={benefit}")
""")

# 2. Create policy with modifier
POST /agent/results/policy-with-modifier
{
  "name": "New benefit for low income",
  "simulation_modifier": "def modify(simulation): ..."
}

# 3. Run economic impact analysis as normal
POST /analysis/economic-impact

Testing

Ran unit tests locally. Integration tests will run on CI.

Stores Python code for structural reforms (custom variable formulas).
Includes database migration.
- Add simulation_modifier field to Policy model (stores Python code)
- Update Modal functions to apply modifiers to simulations
- Add agent upload endpoints for results (deciles, poverty, inequality)
- Update agent system prompt with structural reform workflow
- Database migration for policies.simulation_modifier column

The agent can now write custom variable formulas via simulation modifiers
and run them on UK/US economy simulations.
- Added execute_python tool that runs Python code in a sandbox
- Updated system prompt to instruct agent to test modifiers before submitting
- Agent can now validate syntax and logic before running expensive simulations
Tests for:
- Creating policies with simulation_modifier
- execute_python_code function (basic, error handling, modifier validation)
- PolicyRead including simulation_modifier field
- Added dotenv loading to conftest.py for API keys
- Added test_agent_uses_execute_python_tool - verifies agent can execute Python
- Added test_agent_validates_modifier_code - verifies agent can validate modifiers
@nikhilwoodruff nikhilwoodruff marked this pull request as ready for review January 8, 2026 13:32
@nikhilwoodruff nikhilwoodruff merged commit 405cfdc into main Jan 8, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add structural reform support for AI agent

1 participant