Add Private PE API

integrations/util/cache.py

@@ -1,8 +1,8 @@
-from sentry_sdk import capture_message
+from sentry_sdk import capture_exception
 import datetime
 
 
-class Cache():
+class Cache:
     expire_time = 0
     default = 0
 
@@ -17,8 +17,8 @@ def _update_cache(self):
         try:
             self.data = self.update()
             self.last_update = datetime.datetime.now()
-        except Exception:
-            capture_message(f'Failed to update {self.__class__.__name__}', level='warning')
+        except Exception as e:
+            capture_exception(e, level="warning")
 
     def should_update(self):
         return datetime.datetime.now() > self.last_update + datetime.timedelta(seconds=self.expire_time)

programs/programs/co/pe/member.py

@@ -61,11 +61,12 @@ class Chp(PolicyEngineMembersCalculator):
     def value(self):
         total = 0
 
-        for pkey, pvalue in self.get_data().items():
-            if not self.in_tax_unit(pkey):
+        for member in self.screen.household_members.all():
+            if not self.in_tax_unit(member.id):
                 continue
 
-            if pvalue['co_chp_eligible'][self.pe_period] > 0 and self.screen.has_insurance_types(('none',)):
+            chp_eligible = self.sim.value(self.pe_category, str(member.id), 'co_chp_eligible', self.pe_period) > 0
+            if chp_eligible and self.screen.has_insurance_types(('none',)):
                 total += self.amount
 
         return total

programs/programs/co/pe/tax.py

@@ -38,4 +38,4 @@ def value(self):
                 multiplier = band['percent']
                 break
 
-        return self.get_data()[self.pe_name][self.pe_period] * multiplier
+        return self.get_variable() * multiplier

programs/programs/federal/pe/member.py

@@ -23,9 +23,10 @@ class Wic(PolicyEngineMembersCalculator):
     def value(self):
         total = 0
 
-        for _, pvalue in self.get_data().items():
-            if pvalue[self.pe_name][self.pe_period] > 0:
-                total += self.wic_categories[pvalue['wic_category'][self.pe_period]] * 12
+        for member in self.screen.household_members.all():
+            if self.get_member_variable(member.id) > 0:
+                wic_category = self.sim.value('people', str(member.id), 'wic_category', self.pe_period)
+                total += self.wic_categories[wic_category] * 12
 
         return total
 
@@ -67,18 +68,29 @@ def _value_by_age(self, age: int):
     def value(self):
         total = 0
 
-        for pkey, pvalue in self.get_data().items():
-            # Skip any members that are not in the tax unit.
-            if not self.in_tax_unit(pkey):
-                continue
+        members = self.screen.household_members.all()
 
-            if pvalue[self.pe_name][self.pe_period] <= 0:
+        for member in members:
+            if self.get_member_variable(member.id) <= 0:
                 continue
 
-            total += self._value_by_age(pvalue['age'][self.pe_period])
+            # here we need to adjust for children as policy engine
+            # just uses the average which skews very high for adults and
+            # aged adults
+
+            if self._get_age(member.id) <= 18:
+                medicaid_estimated_value = self.child_medicaid_average
+            elif self._get_age(member.id) > 18 and self._get_age(member.id) < 65:
+                medicaid_estimated_value = self.adult_medicaid_average
+            elif self._get_age(member.id) >= 65:
+                medicaid_estimated_value = self.aged_medicaid_average
+            else:
+                medicaid_estimated_value = 0
+
+            total += medicaid_estimated_value
 
         in_wic_demographic = False
-        for member in self.screen.household_members.all():
+        for member in members:
             if member.pregnant is True or member.age <= 5:
                 in_wic_demographic = True
         if total == 0 and in_wic_demographic:
@@ -89,6 +101,9 @@ def value(self):
 
         return total
 
+    def _get_age(self, member_id: int) -> int:
+        return self.sim.value(self.pe_category, str(member_id), 'age', self.pe_period)
+
 
 class PellGrant(PolicyEngineMembersCalculator):
     pe_name = 'pell_grant'

programs/programs/federal/pe/spm.py

@@ -25,7 +25,7 @@ class Snap(PolicyEngineSpmCalulator):
     pe_output_period = SNAP_PERIOD
 
     def value(self):
-        return int(self.get_data()[self.pe_name][self.pe_output_period]) * 12
+        return int(self.sim.value(self.pe_category, self.pe_sub_category, self.pe_name, self.pe_output_period)) * 12
 
 
 class SchoolLunch(PolicyEngineSpmCalulator):
@@ -39,8 +39,8 @@ def value(self):
         total = 0
         num_children = self.screen.num_children(3, 18)
 
-        if self.get_data()[self.pe_name][self.pe_period] > 0 and num_children > 0:
-            if self.get_data()['school_meal_tier'][self.pe_period] != 'PAID':
+        if self.get_variable() > 0 and num_children > 0:
+            if self.sim.value(self.pe_category, self.pe_sub_category, 'school_meal_tier', self.pe_period) != 'PAID':
                 total = SchoolLunch.amount * num_children
 
         return total

programs/programs/policyengine/calculators/base.py

@@ -4,6 +4,7 @@
 from .dependencies.base import PolicyEngineScreenInput
 from typing import List
 from .constants import YEAR, PREVIOUS_YEAR
+from ..engines import Sim
 
 
 class PolicyEngineCalulator(ProgramCalculator):
@@ -19,9 +20,9 @@ class PolicyEngineCalulator(ProgramCalculator):
     pe_sub_category = ''
     pe_period = YEAR
 
-    def __init__(self, screen: Screen, pe_data):
+    def __init__(self, screen: Screen, sim: Sim):
         self.screen = screen
-        self.pe_data = pe_data
+        self.sim = sim
 
     def eligible(self) -> Eligibility:
         e = Eligibility()
@@ -31,13 +32,13 @@ def eligible(self) -> Eligibility:
         return e
 
     def value(self):
-        return self.get_data()[self.pe_name][self.pe_period]
+        return int(self.get_variable())
 
-    def get_data(self):
+    def get_variable(self):
         '''
-        Return Policy Engine dictionary of the program category and subcategory
+        Return value of the default variable
         '''
-        return self.pe_data[self.pe_category][self.pe_sub_category]
+        return self.sim.value(self.pe_category, self.pe_sub_category, self.pe_name, self.pe_period)
 
     @classmethod
     def can_calc(cls, missing_dependencies: Dependencies):
@@ -47,37 +48,40 @@ def can_calc(cls, missing_dependencies: Dependencies):
 
         return True
 
+
 class PolicyEngineTaxUnitCalulator(PolicyEngineCalulator):
     pe_category = 'tax_units'
     pe_sub_category = 'tax_unit'
     tax_unit_dependent = True
     pe_period = PREVIOUS_YEAR
 
+
 class PolicyEngineSpmCalulator(PolicyEngineCalulator):
     pe_category = 'spm_units'
     pe_sub_category = 'spm_unit'
 
+
 class PolicyEngineMembersCalculator(PolicyEngineCalulator):
     tax_unit_dependent = True
     pe_category = 'people'
 
     def value(self):
         total = 0
-        for pkey, pvalue in self.get_data().items():
+        for member in self.screen.household_members.all():
             # The following programs use income from the tax unit,
             # so we want to skip any members that are not in the tax unit.
-            if not self.in_tax_unit(pkey) and self.tax_unit_dependent:
+            if not self.in_tax_unit(member.id) and self.tax_unit_dependent:
                 continue
 
-            pe_value = pvalue[self.pe_name][self.pe_period]
+            pe_value = self.get_member_variable(member.id)
 
             total += pe_value
 
         return total
 
-    def in_tax_unit(self, member_id) -> bool:
-        return str(member_id) in self.pe_data['tax_units']['tax_unit']['members']
+    def in_tax_unit(self, member_id: int) -> bool:
+        return str(member_id) in self.sim.members('tax_units', 'tax_unit')
 
-    def get_data(self):
-        return self.pe_data[self.pe_category]
+    def get_member_variable(self, member_id: int):
+        return self.sim.value(self.pe_category, str(member_id), self.pe_name, self.pe_period)
 

programs/programs/policyengine/engines.py

@@ -0,0 +1,114 @@
+from integrations.util.cache import Cache
+from decouple import config
+import requests
+
+
+class Sim:
+    method_name = ""
+
+    def __init__(self, data) -> None:
+        self.data = data
+
+    def value(self, unit, sub_unit, variable, period):
+        """
+        Calculate variable at the period
+        """
+        raise NotImplementedError
+
+    def members(self, unit, sub_unit):
+        """
+        Return a list of the members in the sub unit
+        """
+        raise NotImplementedError
+
+
+class ApiSim(Sim):
+    method_name = "Policy Engine API"
+    pe_url = "https://api.policyengine.org/us/calculate"
+
+    def __init__(self, data) -> None:
+        response = requests.post(self.pe_url, json=data)
+        self.data = response.json()["result"]
+
+    def value(self, unit, sub_unit, variable, period):
+        return self.data[unit][sub_unit][variable][period]
+
+    def members(self, unit, sub_unit):
+        return self.data[unit][sub_unit]["members"]
+
+
+class PolicyEngineBearerTokenCache(Cache):
+    expire_time = 60 * 60 * 24 * 29
+    default = ""
+    client_id: str = config("POLICY_ENGINE_CLIENT_ID", "")
+    client_secret: str = config("POLICY_ENGINE_CLIENT_SECRET", "")
+    domain = "https://policyengine.uk.auth0.com"
+    endpoint = "/oauth/token"
+
+    def update(self):
+        # https://policyengine.org/us/api#fetch_token
+
+        if self.client_id == "" or self.client_secret == "":
+            raise Exception("client id or secret not configured")
+
+        payload = {
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "grant_type": "client_credentials",
+            "audience": "https://household.api.policyengine.org",
+        }
+
+        res = requests.post(self.domain + self.endpoint, json=payload)
+
+        return res.json()["access_token"]
+
+
+class PrivateApiSim(ApiSim):
+    method_name = "Private Policy Engine API"
+    token = PolicyEngineBearerTokenCache()
+    pe_url = "https://household.api.policyengine.org/us/calculate"
+
+    def __init__(self, data) -> None:
+        token = self.token.fetch()
+
+        headers = {
+            "Authorization": f"Bearer {token}",
+        }
+
+        res = requests.post(self.pe_url, json=data, headers=headers)
+
+        self.data = res.json()["result"]
+
+
+# NOTE: Code to run Policy Engine locally. This is currently too CPU expensive to run in production.
+# Requires the Policy Engine package to be installed and imported.
+#
+# class LocalSim(Sim):
+#     method_name = 'local package'
+#
+#     def __init__(self, data) -> None:
+#         self.household = data['household']
+#
+#         self.entity_map = {}
+#         for entity in self.household.keys():
+#             group_map = {}
+#
+#             for i, group in enumerate(self.household[entity].keys()):
+#                 group_map[group] = i
+#
+#             self.entity_map[entity] = group_map
+#
+#         self.sim = Simulation(situation=self.household)
+#
+#     def value(self, unit, sub_unit, variable, period):
+#         data = self.sim.calculate(variable, period)
+#
+#         index = self.entity_map[unit][sub_unit]
+#
+#         return data[index]
+#
+#     def members(self, unit, sub_unit):
+#         return self.household[unit][sub_unit]['members']
+
+
+pe_engines: list[Sim] = [PrivateApiSim, ApiSim]