from pydantic import BaseModel, Field, BeforeValidator, model_validator from typing import Annotated, List, Optional, Literal # Example constants for validation TYPICAL_MEASURE_TYPES = [ "wall_insulation", "roof_insulation", "ventilation", "floor_insulation", "windows", "fireplace", "heating", "hot_water", "low_energy_lighting", "secondary_heating", "solar_pv" ] WALL_INSULATION_MEASURES = ["internal_wall_insulation", "external_wall_insulation", "cavity_wall_insulation"] ROOF_INSULATION_MEASURES = [ "loft_insulation", "flat_roof_insulation", "room_roof_insulation", "sloping_ceiling_insulation" ] WALL_INSULATION_WITH_VENTILATION_MEASURES = [ "internal_wall_insulation+mechanical_ventilation", "external_wall_insulation+mechanical_ventilation", "cavity_wall_insulation+mechanical_ventilation" ] # Both all and roof insulaiton measures are eligible for ECO4. These are the remaining fabric and heating measures # This is based on th measures we have recommendations for ECO4_ELIGIBILE_FABRIC_MEASURES = [ "suspended_floor_insulation", "solid_floor_insulation", "double_glazing", "secondary_glazing" ] ECO4_ELIGIBLE_HEATING_MEASURES = [ "boiler_upgrade", "high_heat_retention_storage_heaters", "air_source_heat_pump", "solar_pv" ] SPECIFIC_MEASURES = ( WALL_INSULATION_MEASURES + ROOF_INSULATION_MEASURES + ECO4_ELIGIBILE_FABRIC_MEASURES + ECO4_ELIGIBLE_HEATING_MEASURES + [ "secondary_heating", "ventilation", "low_energy_lighting", "fireplace", "hot_water_tank_insulation", "cylinder_thermostat" ] ) INSULATION_MEASURES = [ "internal_wall_insulation", "external_wall_insulation", "cavity_wall_insulation", "loft_insulation", "flat_roof_insulation", "room_roof_insulation", "sloping_ceiling_insulation", "suspended_floor_insulation", "solid_floor_insulation", ] NON_INVASIVE_SPECIFIC_MEASURES = [ "trickle_vents", "draught_proofing", "mixed_glazing", "cavity_extract_and_refill", "extension_cavity_wall_insulation" ] # This allows us to extend high level categories for measures such as "wall_insulation" to the specific measures # such as "external_wall_insulation", "internal_wall_insulation", "cavity_wall_insulation" MEASURE_MAP = { "wall_insulation": [ "internal_wall_insulation", "external_wall_insulation", "cavity_wall_insulation", ], "roof_insulation": [ "loft_insulation", "flat_roof_insulation", "room_roof_insulation", "sloping_ceiling_insulation" ], "floor_insulation": ["suspended_floor_insulation", "solid_floor_insulation"], "heating": ["boiler_upgrade", "high_heat_retention_storage_heaters", "air_source_heat_pump"], "windows": ["double_glazing", "secondary_glazing"], "heating_controls": ["roomstat_programmer_trvs", "time_temperature_zone_control"] } VALID_GOALS = ["Increasing EPC", "Energy Savings", "Reducing CO2 emissions"] VALID_HOUSING_TYPES = ["Social", "Private"] VALID_EVENT_TYPES = ["remote_assessment", "eco_project"] # Define the validation function for inclusions/exclusions def check_inclusion_or_exclusion(value: str) -> str: if value not in TYPICAL_MEASURE_TYPES + SPECIFIC_MEASURES + NON_INVASIVE_SPECIFIC_MEASURES: raise ValueError(f"{value} is not an allowed inclusion") return value def check_goals(value: str) -> str: assert value in VALID_GOALS, f"{value} is not a valid goal" return value def check_housing_type(value: str) -> str: assert value in VALID_HOUSING_TYPES, f"{value} is not a valid housing type" return value def check_event_type(value: str) -> str: assert value in VALID_EVENT_TYPES, f"{value} is not a valid event type" return value # Use Annotated with BeforeValidator for each list item validation InclusionOrExclusionItem = Annotated[str, BeforeValidator(check_inclusion_or_exclusion)] Goal = Annotated[str, BeforeValidator(check_goals)] HousingType = Annotated[str, BeforeValidator(check_housing_type)] EventType = Annotated[str, BeforeValidator(check_event_type)] class PlanTriggerRequest(BaseModel): budget: Optional[float] = None goal: Goal housing_type: HousingType goal_value: Optional[str] = None portfolio_id: int trigger_file_path: str already_installed_file_path: Optional[str] = None patches_file_path: Optional[str] = None non_invasive_recommendations_file_path: Optional[str] = None valuation_file_path: Optional[str] = None exclusions: Optional[List[InclusionOrExclusionItem]] = Field(default=None, min_length=0) inclusions: Optional[List[InclusionOrExclusionItem]] = Field(default=None, min_length=0) # This is a list of measures that we want to be included, if they are options # Default to empty required_measures: Optional[List[InclusionOrExclusionItem]] = Field(default=[], min_length=0) scenario_name: Optional[str] = "" scenario_id: Optional[str | int] = None # Used to utilise and existing scenario for a engine run multi_plan: Optional[bool] = False default_u_values: Optional[bool] = True ashp_cop: Optional[float] = 2.8 # When performing a remote assessment, if this has been set, it will allow the engine to # pull data from the find my epc website, to utilise as part of a remote assessment event_type: Optional[Literal["remote_assessment", "eco_project"]] = None # If true, before optimising the engine will select a slightly larger package, to account for the SAP 10 causing # scores to drop by a few points simulate_sap_10: Optional[bool] = False # Add in optional fields which describe the format of the asset list being used file_type: Optional[Literal["csv", "xlsx"]] = None file_format: Optional[Literal["domna_asset_list", "ara_property_list"]] = None sheet_name: Optional[str] = None sheet_count: Optional[int] = None # If one of index_start or index_end is set, the other must be set too index_start: Optional[int] = None index_end: Optional[int] = None # Task and subtask IDs task_id: Optional[str] = None subtask_id: Optional[str] = None # Optional flag to trigger a fabric first task enforce_fabric_first: Optional[bool] = False @model_validator(mode="after") def check_indexes(self): if (self.index_start is None) != (self.index_end is None): raise ValueError("Both index_start and index_end must be set or both must be None") return self @model_validator(mode="after") def check_goal_value_requirement(self): # Make sure that goal_value is set when goal is "Increasing EPC" if self.goal == "Increasing EPC" and not self.goal_value: raise ValueError("goal_value is required when goal is 'Increasing EPC'") return self