Merge pull request #257 from PolicyEngine/fix/allow-multiple-group-entities

Allow multiple stored household group entities

Author: anth-volk

changelog.d/256.fixed.md

@@ -0,0 +1 @@
+Stored `/households` definitions now use the same plural entity-list contract as household calculation, including validation for multi-group person/entity linkage and compatibility coercion for legacy singular payloads.

src/policyengine_api/api/household.py

@@ -96,9 +96,6 @@ class HouseholdCalculateRequest(BaseModel):
     }
     """
 
-    country_id: CountryId = Field(
-        description="Which country model to use ('us' or 'uk')"
-    )
     people: list[dict[str, Any]] = Field(
         description="List of people with flat variable values. Include person_id and person_{entity}_id fields to link to entities."
     )
@@ -130,6 +127,9 @@ class HouseholdCalculateRequest(BaseModel):
         default=None,
         description="Simulation year (default: 2024 for US, 2026 for UK). Specify this instead of embedding years in variable values.",
     )
+    country_id: CountryId = Field(
+        description="Which country model to use ('us' or 'uk')"
+    )
     policy_id: UUID | None = Field(
         default=None, description="Optional policy reform ID"
     )
@@ -183,9 +183,6 @@ class HouseholdImpactRequest(BaseModel):
     }
     """
 
-    country_id: CountryId = Field(
-        description="Which country model to use ('us' or 'uk')"
-    )
     people: list[dict[str, Any]] = Field(
         description="List of people with flat variable values. Include person_id and person_{entity}_id fields to link to entities."
     )
@@ -216,6 +213,9 @@ class HouseholdImpactRequest(BaseModel):
     year: int | None = Field(
         default=None, description="Simulation year (default: 2024 for US, 2026 for UK)"
     )
+    country_id: CountryId = Field(
+        description="Which country model to use ('us' or 'uk')"
+    )
     policy_id: UUID | None = Field(
         default=None, description="Reform policy ID to compare against baseline"
     )
@@ -854,7 +854,7 @@ def calculate_household(
     Use flat values for all variables - do NOT use time-period format like {"2024": value}.
     The simulation year is specified via the `year` parameter.
 
-    US example: people=[{"employment_income": 70000, "age": 40}], tax_unit={"state_code": "CA"}, year=2024
+    US example: people=[{"employment_income": 70000, "age": 40}], tax_unit=[{"state_code": "CA"}], year=2024
     UK example: people=[{"employment_income": 50000, "age": 30}], year=2026
     """
     with logfire.span(

src/policyengine_api/api/household_analysis.py

@@ -9,6 +9,10 @@
 3. Run analysis: POST /analysis/household-impact (returns report_id)
 4. Poll GET /analysis/household-impact/{report_id} until status="completed"
 5. Results include baseline_result, reform_result (if comparison), and impact diff
+
+Stored households now use the same plural entity-list contract as the
+/household/calculate endpoint. Analysis should pass those lists through
+without collapsing multi-group households into a single entity row.
 """
 
 from dataclasses import dataclass
@@ -109,7 +113,7 @@ def calculate_uk_household(
     year: int,
     policy_data: dict | None,
 ) -> dict:
-    """Calculate UK household using the existing implementation."""
+    """Calculate UK household using the stored-household plural entity contract."""
     from policyengine_api.api.household import _calculate_household_uk
 
     return _calculate_household_uk(
@@ -126,7 +130,7 @@ def calculate_us_household(
     year: int,
     policy_data: dict | None,
 ) -> dict:
-    """Calculate US household using the existing implementation."""
+    """Calculate US household using the stored-household plural entity contract."""
     from policyengine_api.api.household import _calculate_household_us
 
     return _calculate_household_us(

src/policyengine_api/api/households.py

@@ -34,9 +34,7 @@ def _pack_household_data(body: HouseholdCreate) -> dict[str, Any]:
     """Pack the flat request fields into a single JSON blob for storage."""
     data: dict[str, Any] = {"people": body.people}
     for key in _ENTITY_GROUP_KEYS:
-        val = getattr(body, key)
-        if val is not None:
-            data[key] = val
+        data[key] = list(getattr(body, key))
     return data
 
 
@@ -49,12 +47,12 @@ def _to_read(record: Household) -> HouseholdRead:
         year=record.year,
         label=record.label,
         people=data["people"],
-        tax_unit=data.get("tax_unit"),
-        family=data.get("family"),
-        spm_unit=data.get("spm_unit"),
-        marital_unit=data.get("marital_unit"),
-        household=data.get("household"),
-        benunit=data.get("benunit"),
+        tax_unit=data.get("tax_unit", []),
+        family=data.get("family", []),
+        spm_unit=data.get("spm_unit", []),
+        marital_unit=data.get("marital_unit", []),
+        household=data.get("household", []),
+        benunit=data.get("benunit", []),
         created_at=record.created_at,
         updated_at=record.updated_at,
     )

src/policyengine_api/models/__init__.py

@@ -34,6 +34,7 @@
     HouseholdJobRead,
     HouseholdJobStatus,
 )
+from .household_payload import HouseholdEntityCollections, StoredHouseholdPayload
 from .inequality import Inequality, InequalityCreate, InequalityRead
 from .intra_decile_impact import (
     DecileType,
@@ -152,6 +153,8 @@
     "DynamicRead",
     "Household",
     "HouseholdCreate",
+    "HouseholdEntityCollections",
+    "StoredHouseholdPayload",
     "HouseholdRead",
     "HouseholdJob",
     "HouseholdJobCreate",

src/policyengine_api/models/household.py

@@ -4,10 +4,32 @@
 from typing import Any
 from uuid import UUID, uuid4
 
+from pydantic import field_validator, model_validator
 from sqlalchemy import JSON
 from sqlmodel import Column, Field, SQLModel
 
-from policyengine_api.config.constants import CountryId
+from policyengine_api.models.household_payload import StoredHouseholdPayload
+
+_ENTITY_ID_KEY_BY_GROUP = {
+    "benunit": "benunit_id",
+    "marital_unit": "marital_unit_id",
+    "family": "family_id",
+    "spm_unit": "spm_unit_id",
+    "tax_unit": "tax_unit_id",
+    "household": "household_id",
+}
+
+_ENTITY_GROUP_FIELDS = tuple(_ENTITY_ID_KEY_BY_GROUP.keys())
+
+
+def _coerce_entity_group_collection(value: Any) -> Any:
+    if value is None:
+        return []
+    if isinstance(value, list):
+        return value
+    if isinstance(value, dict):
+        return [value]
+    return value
 
 
 class HouseholdBase(SQLModel):
@@ -29,23 +51,70 @@ class Household(HouseholdBase, table=True):
     updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
 
 
-class HouseholdCreate(SQLModel):
+class HouseholdCreate(StoredHouseholdPayload):
     """Schema for creating a stored household.
 
-    Accepts the flat structure matching the frontend Household interface:
-    people as an array, entity groups as optional dicts.
+    Uses the same list-based relational entity shape as the household
+    calculation APIs.
     """
 
-    country_id: CountryId
-    year: int
-    label: str | None = None
-    people: list[dict[str, Any]]
-    tax_unit: dict[str, Any] | None = None
-    family: dict[str, Any] | None = None
-    spm_unit: dict[str, Any] | None = None
-    marital_unit: dict[str, Any] | None = None
-    household: dict[str, Any] | None = None
-    benunit: dict[str, Any] | None = None
+    @field_validator(*_ENTITY_GROUP_FIELDS, mode="before")
+    @classmethod
+    def coerce_legacy_singular_entity_groups(cls, value: Any) -> Any:
+        return _coerce_entity_group_collection(value)
+
+    @model_validator(mode="after")
+    def validate_relationships(self) -> "HouseholdCreate":
+        person_ids = [
+            person["person_id"]
+            for person in self.people
+            if person.get("person_id") is not None
+        ]
+        if len(person_ids) != len(set(person_ids)):
+            raise ValueError("people contains duplicate person_id values")
+
+        for group_key, entity_id_key in _ENTITY_ID_KEY_BY_GROUP.items():
+            entity_records = getattr(self, group_key)
+            person_link_key = f"person_{entity_id_key}"
+
+            entity_ids = [
+                entity[entity_id_key]
+                for entity in entity_records
+                if entity.get(entity_id_key) is not None
+            ]
+
+            if len(entity_ids) != len(set(entity_ids)):
+                raise ValueError(
+                    f"{group_key} contains duplicate {entity_id_key} values"
+                )
+
+            requires_linkage = len(entity_records) > 1
+            person_links = []
+            for person in self.people:
+                person_link = person.get(person_link_key)
+                if person_link is None:
+                    if requires_linkage:
+                        raise ValueError(
+                            f"people must include {person_link_key} when {group_key} has multiple rows"
+                        )
+                    continue
+                person_links.append(person_link)
+
+            if not person_links:
+                continue
+
+            if len(entity_ids) != len(entity_records):
+                raise ValueError(
+                    f"{group_key} rows must all include {entity_id_key} when people reference {person_link_key}"
+                )
+
+            unknown_links = sorted(set(person_links) - set(entity_ids))
+            if unknown_links:
+                raise ValueError(
+                    f"{group_key} is missing rows for referenced {entity_id_key} values: {unknown_links}"
+                )
+
+        return self
 
 
 class HouseholdRead(HouseholdCreate):

src/policyengine_api/models/household_payload.py

@@ -0,0 +1,27 @@
+"""Shared household payload models."""
+
+from typing import Any
+
+from sqlmodel import Field, SQLModel
+
+from policyengine_api.config.constants import CountryId
+
+
+class HouseholdEntityCollections(SQLModel):
+    """Plural household entity collections used by stored and calculation APIs."""
+
+    benunit: list[dict[str, Any]] = Field(default_factory=list)
+    marital_unit: list[dict[str, Any]] = Field(default_factory=list)
+    family: list[dict[str, Any]] = Field(default_factory=list)
+    spm_unit: list[dict[str, Any]] = Field(default_factory=list)
+    tax_unit: list[dict[str, Any]] = Field(default_factory=list)
+    household: list[dict[str, Any]] = Field(default_factory=list)
+
+
+class StoredHouseholdPayload(HouseholdEntityCollections):
+    """Core payload shared by stored household create/read flows."""
+
+    country_id: CountryId
+    people: list[dict[str, Any]]
+    year: int
+    label: str | None = None

test_fixtures/fixtures_households.py

@@ -14,9 +14,9 @@
         {"age": 30, "employment_income": 50000},
         {"age": 28, "employment_income": 30000},
     ],
-    "tax_unit": {},
-    "family": {},
-    "household": {"state_name": "CA"},
+    "tax_unit": [{}],
+    "family": [{}],
+    "household": [{"state_name": "CA"}],
 }
 
 MOCK_UK_HOUSEHOLD_CREATE = {
@@ -26,8 +26,88 @@
     "people": [
         {"age": 40, "employment_income": 35000},
     ],
-    "benunit": {"is_married": False},
-    "household": {"region": "LONDON"},
+    "benunit": [{"is_married": False}],
+    "household": [{"region": "LONDON"}],
+}
+
+MOCK_US_MULTI_GROUP_HOUSEHOLD_CREATE = {
+    "country_id": "us",
+    "year": 2024,
+    "label": "US multi-group household",
+    "people": [
+        {
+            "person_id": 0,
+            "person_household_id": 0,
+            "person_tax_unit_id": 0,
+            "person_marital_unit_id": 0,
+            "age": 30,
+            "employment_income": 50000,
+        },
+        {
+            "person_id": 1,
+            "person_household_id": 0,
+            "person_tax_unit_id": 0,
+            "person_marital_unit_id": 1,
+            "age": 28,
+            "employment_income": 30000,
+        },
+    ],
+    "tax_unit": [{"tax_unit_id": 0, "state_name": "CA"}],
+    "marital_unit": [
+        {"marital_unit_id": 0},
+        {"marital_unit_id": 1},
+    ],
+    "family": [{"family_id": 0}],
+    "spm_unit": [{"spm_unit_id": 0}],
+    "household": [{"household_id": 0, "state_name": "CA"}],
+}
+
+MOCK_US_FULL_MULTI_GROUP_HOUSEHOLD_CREATE = {
+    "country_id": "us",
+    "year": 2024,
+    "label": "US fully multi-group household",
+    "people": [
+        {
+            "person_id": 0,
+            "person_household_id": 0,
+            "person_tax_unit_id": 0,
+            "person_marital_unit_id": 0,
+            "person_family_id": 0,
+            "person_spm_unit_id": 0,
+            "age": 30,
+            "employment_income": 50000,
+        },
+        {
+            "person_id": 1,
+            "person_household_id": 1,
+            "person_tax_unit_id": 1,
+            "person_marital_unit_id": 1,
+            "person_family_id": 1,
+            "person_spm_unit_id": 1,
+            "age": 28,
+            "employment_income": 30000,
+        },
+    ],
+    "tax_unit": [
+        {"tax_unit_id": 0, "state_name": "CA"},
+        {"tax_unit_id": 1, "state_name": "CA"},
+    ],
+    "marital_unit": [
+        {"marital_unit_id": 0},
+        {"marital_unit_id": 1},
+    ],
+    "family": [
+        {"family_id": 0},
+        {"family_id": 1},
+    ],
+    "spm_unit": [
+        {"spm_unit_id": 0},
+        {"spm_unit_id": 1},
+    ],
+    "household": [
+        {"household_id": 0, "state_name": "CA"},
+        {"household_id": 1, "state_name": "NY"},
+    ],
 }
 
 MOCK_HOUSEHOLD_MINIMAL = {
@@ -36,6 +116,19 @@
     "people": [{"age": 25}],
 }
 
+MOCK_US_HOUSEHOLD_CREATE_LEGACY = {
+    "country_id": "us",
+    "year": 2024,
+    "label": "US legacy household",
+    "people": [
+        {"age": 30, "employment_income": 50000},
+        {"age": 28, "employment_income": 30000},
+    ],
+    "tax_unit": {},
+    "family": {},
+    "household": {"state_name": "CA"},
+}
+
 
 # -----------------------------------------------------------------------------
 # Factory functions