diff --git a/backend/app/db/functions/recommendations_functions.py b/backend/app/db/functions/recommendations_functions.py index aa966fbb..d4c3fcb9 100644 --- a/backend/app/db/functions/recommendations_functions.py +++ b/backend/app/db/functions/recommendations_functions.py @@ -625,6 +625,13 @@ def get_plans_by_portfolio_id(portfolio_id: int) -> List[PlanModel]: return session_any.exec(stmt).scalars().all() +def get_plans_by_ids(ids: List[int]) -> List[PlanModel]: + stmt = select(PlanModel).where(PlanModel.id.in_(ids)) + with db_read_session() as session: + session_any: Any = session # Typehint as Any to satisfy Pylance... + return session_any.exec(stmt).scalars().all() + + def get_scenarios_by_portfolio_id(portfolio_id: int) -> List[ScenarioModel]: stmt = select(ScenarioModel).where(ScenarioModel.portfolio_id == portfolio_id) with db_read_session() as session: diff --git a/backend/categorisation/categorisation_trigger_request.py b/backend/categorisation/categorisation_trigger_request.py index aa2b8ed3..46ce6f1c 100644 --- a/backend/categorisation/categorisation_trigger_request.py +++ b/backend/categorisation/categorisation_trigger_request.py @@ -5,4 +5,5 @@ from pydantic import BaseModel class CategorisationTriggerRequest(BaseModel): portfolio_id: int - plan_priority_order: Optional[List[int]] + plans_to_consider: Optional[List[int]] = None + plan_priority_order: Optional[List[int]] = None diff --git a/backend/categorisation/handler/handler.py b/backend/categorisation/handler/handler.py index 20076613..449c5ccf 100644 --- a/backend/categorisation/handler/handler.py +++ b/backend/categorisation/handler/handler.py @@ -20,7 +20,11 @@ def handler(event: Mapping[str, Any], context: Any) -> None: logger.debug("Successfully validated request body") - process_portfolio(payload.portfolio_id) + process_portfolio( + payload.portfolio_id, + payload.plans_to_consider, + payload.plan_priority_order, + ) except Exception as e: logger.error(f"Failed to process record: {e}") diff --git a/backend/categorisation/processor.py b/backend/categorisation/processor.py index 184ccac2..b7ddfc62 100644 --- a/backend/categorisation/processor.py +++ b/backend/categorisation/processor.py @@ -1,10 +1,9 @@ from collections import defaultdict from typing import Dict, List, Optional -from sqlalchemy import Tuple - from backend.app.db.functions.recommendations_functions import ( bulk_update_plans, + get_plans_by_ids, get_plans_by_portfolio_id, get_scenarios_by_portfolio_id, ) @@ -17,10 +16,14 @@ logger = setup_logger() def process_portfolio( - portfolio_id: int, plan_priority_order: Optional[List[int]] = [] + portfolio_id: int, + plans_to_consider: Optional[List[int]] = None, + plan_priority_order: Optional[List[int]] = None, ) -> None: logger.info(f"Processing portfolio {portfolio_id}") - plans: List[Plan] = _load_plans_for_portfolio(portfolio_id) + + plans: List[Plan] = _load_plans_for_portfolio(portfolio_id, plans_to_consider) + plans_by_property: Dict[int, List[Plan]] = _group_plans_by_property(plans) for uprn, property_plans in plans_by_property.items(): @@ -74,10 +77,25 @@ def choose_cheapest_relevant_plan( return cheapest_plans[0] -def _load_plans_for_portfolio(portfolio_id: int) -> List[Plan]: +def _load_plans_for_portfolio( + portfolio_id: int, plans_to_consider: Optional[List[int]] = None +) -> List[Plan]: + + if plans_to_consider: + if len(plans_to_consider) < 2: + raise ValueError("Cannot run auto categorisation for fewer than 2 plans") + + logger.info(f"Getting {len(plans_to_consider)} Plans") + plan_models: List[PlanModel] = get_plans_by_ids(plans_to_consider) + + else: + logger.info( + f"No list of Plans to consider provided. Getting all Plans for portfolio {portfolio_id}" + ) + plan_models: List[PlanModel] = get_plans_by_portfolio_id(portfolio_id) + plans: List[Plan] = [] - plan_models = get_plans_by_portfolio_id(portfolio_id) scenarios: List[ScenarioModel] = get_scenarios_by_portfolio_id(portfolio_id) if not scenarios: