option to only consider a specific list of plans

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:

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

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}")

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: