Batch planning

PyAres/Models/__init__.py

@@ -1,4 +1,4 @@
-from .ares_data_models import AresDataType, Outcome, RequestMetadata, AresSchemaEntry, Quantity, QuantitySchema, Limits
+from .ares_data_models import AresDataType, Outcome, RequestMetadata, AresSchemaEntry, Quantity, QuantitySchema, Limits, PlanStatusCode
 
 __all__ = [
     "AresDataType",
@@ -7,5 +7,6 @@
     "AresSchemaEntry",
     "Quantity",
     "QuantitySchema",
-    "Limits"
+    "Limits",
+    "PlanStatusCode"
 ]

PyAres/Models/ares_data_models.py

@@ -30,6 +30,12 @@ class Outcome(Enum):
     WARNING = 3
     CANCELED = 4
 
+class PlanStatusCode(Enum):
+    PLAN_STATUS_UNSPECIFIED = 0
+    PLAN_ACCEPTED = 1
+    PLAN_UNACHIEVABLE = 2
+    PLAN_FAILED = 3
+
 class RequestMetadata():
     def __init__(self, proto_metadata: request_metadata_pb2.RequestMetadata):
         self.system_name = proto_metadata.system_name

PyAres/Planning/__init__.py

@@ -1,9 +1,12 @@
-from .planner_models import PlanningParameter, PlanRequest, PlanResponse, ParameterHistoryItem
+from .planner_models import PlanningParameter, PlanRequest, PlanResponse, ParameterHistoryItem, Plan, PlannedParameter
 from .planning_service import AresPlannerService
 
 __all__ = [
     "PlanningParameter",
     "PlanRequest",
     "PlanResponse",
-    "AresPlannerService"
+    "AresPlannerService",
+    "Plan",
+    "PlannedParameter",
+    "ParameterHistoryItem",
 ]

PyAres/Planning/planner_models.py

@@ -1,5 +1,6 @@
 from typing import Dict, Any, List, Sequence, Optional
-from ..Models import Outcome, AresDataType, RequestMetadata
+from ..Models import Outcome, AresDataType, RequestMetadata, PlanStatusCode
+from enum import Enum
 
 class ParameterHistoryItem:
     """ Represents a single historical parameter item """
@@ -107,7 +108,13 @@ class PlanRequest:
 
     Designed to provide a more user-friendly abstraction for interacting with a plan request message.
     """
-    def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any], analysis_results: Sequence[float], metadata: RequestMetadata = RequestMetadata.from_default_values()):
+    def __init__(self, 
+                parameters: list[PlanningParameter], 
+                settings: Dict[str, Any], 
+                analysis_results: Sequence[float], 
+                batch_size: int = 1,
+                metadata: RequestMetadata = RequestMetadata.from_default_values(),
+                previous_plan_status_codes: List[PlanStatusCode] = []):
         """
         Initializes a PlanRequest.
 
@@ -117,7 +124,9 @@ def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any]
         self.parameters = parameters
         self.settings = settings
         self.analysis_results = analysis_results
+        self.batch_size = batch_size
         self.request_metadata = metadata
+        self.previous_plan_status_codes = previous_plan_status_codes
 
     def __str__(self) -> str:
         param_str = "\n ".join(self.parameter_names)
@@ -127,13 +136,15 @@ def __str__(self) -> str:
         metadata_str = str(self.request_metadata).replace('\n', '\n ')
         return (f"PlanRequest object with:\n"
                 f"parameters:\n"
-                f" {param_str}\n"
+                f"{param_str}\n"
                 f"settings:\n"
-                f" {settings_str}\n"
+                f"{settings_str}\n"
                 f"analysis_results:\n"
-                f" {analysis_str}\n"
+                f"{analysis_str}\n"
                 f"request_metadata:\n"
-                f"{metadata_str}")
+                f"{metadata_str}"
+                f"batch_size:\n"
+                f"{self.batch_size}")
 
     def __repr__(self) -> str:
         return self.__str__()
@@ -175,6 +186,8 @@ def __init__(self,
             parameter_names: A list of names associated with planned parameters.
             parameter_values: A list of values associated with planned parameters. 
             parameter_data: A python dictionary of key:value pairs of planned parameters and planned values
+            outcome: An enum of type Outcome that determines whether the planning process succeeded or not, defaults to SUCCESS
+            error_string: An optional string for specifying planning failure reasons to be relayed to ARES
         """
         if parameter_data is not None:
             self.parameter_names = list(parameter_data.keys())
@@ -200,4 +213,39 @@ def __str__(self):
                 f" error_string: {self.error_string}\n")
 
     def __repr__(self) -> str:
-        return self.__str__()
+        return self.__str__()
+
+class PlannedParameter:
+    def __init__(self, parameter_name: str, parameter_value: Any):
+        self.parameter_name = parameter_name
+        self.parameter_value = parameter_value
+
+    def __str__(self):
+        # A clean, readable key-value output
+        return f"{self.parameter_name}: {self.parameter_value}"
+
+    def __repr__(self):
+        # The !r formatting flag automatically wraps strings in quotes and calls __repr__ on the values
+        return f"PlannedParameters(parameter_name={self.parameter_name!r}, parameter_value={self.parameter_value!r})"
+
+class Plan:
+    def __init__(self, planned_parameters: List[PlannedParameter], outcome: Outcome, error_string: str = ""):
+        self.planned_parameters = planned_parameters
+        self.outcome = outcome
+        self.error_string = error_string
+
+    def __str__(self):
+        base_str = f"Plan (Outcome: {self.outcome})"
+
+        # Format the list of parameters into a readable string
+        if self.planned_parameters:
+            params_str = ", ".join(str(p) for p in self.planned_parameters)
+            base_str += f" | Parameters: [{params_str}]"
+
+        if self.error_string:
+            base_str += f" - Error: '{self.error_string}'"
+
+        return base_str
+
+    def __repr__(self):
+        return f"Plan(planned_parameters={self.planned_parameters!r}, outcome={self.outcome!r}, error_string={self.error_string!r})"

PyAres/Planning/planning_service.py

@@ -21,13 +21,15 @@
 from ..Utils import ares_data_type_utils
 from ..Utils import ares_struct_utils
 from ..Utils import ares_outcome_utils
+from ..Utils import ares_plan_status_code_utils
+from ..Utils import plan_response_utils
 
 # Import python models
 from ..Models import ares_data_models, Limits
 from .planner_models import *
 
 # Type hint for the user's custom planning logic
-PlanLogicFunction = Callable[[PlanRequest], Union[PlanResponse, Awaitable[PlanResponse]]]
+PlanLogicFunction = Callable[[PlanRequest], Union[PlanResponse, Awaitable[PlanResponse], List[Plan], Awaitable[List[Plan]]]]
 
 class AresPlannerServiceWrapper(planner_service_grpc.AresRemotePlannerServiceServicer):
     """
@@ -122,7 +124,9 @@ def Plan(self, request: plan_pb2.PlanningRequest, context) -> plan_pb2.PlanningR
         python_request = PlanRequest(parameters=parameters, 
                                      settings=ares_struct_utils.ares_struct_to_dict(request.adapter_settings), 
                                      analysis_results=list(request.analysis_results),
-                                     metadata=RequestMetadata(request.metadata))
+                                     metadata=RequestMetadata(request.metadata),
+                                     batch_size=request.batch_size,
+                                     previous_plan_status_codes=[ares_plan_status_code_utils.proto_plan_status_to_python_plan_status(c) for c in request.previous_plan_status_codes])
 
         #Handle call using the user's custom planning logic 
         response_proto = plan_pb2.PlanningResponse()
@@ -139,20 +143,23 @@ def Plan(self, request: plan_pb2.PlanningRequest, context) -> plan_pb2.PlanningR
             response_proto.planning_outcome = ares_outcome_enum_pb2.FAILURE
             return response_proto
 
-        if not isinstance(python_response, PlanResponse):
-            response_proto.error_string = "The returned response from the user planning method was not a plan response, and thus was invalid."
-            response_proto.planning_outcome = ares_outcome_enum_pb2.FAILURE
-            return response_proto
-
-        response_proto.planning_outcome = ares_outcome_utils.python_ares_outcome_to_proto_ares_outcome(python_response.outcome)
-        response_proto.error_string = python_response.error_string
+        if isinstance(python_response, PlanResponse):
+            response_proto.planning_outcome = ares_outcome_utils.python_ares_outcome_to_proto_ares_outcome(python_response.outcome)
+            response_proto.error_string = python_response.error_string
+
+            for i in range(len(python_response.parameter_names)):
+                planned_parameter = plan_pb2.PlannedParameter(parameter_value=ares_value_utils.create_ares_value(python_response.parameter_values[i]))
+                planned_parameter.parameter_name = python_response.parameter_names[i]
+                new_planned_parameter = response_proto.planned_parameters.add()
+                new_planned_parameter.CopyFrom(planned_parameter)
 
-        for i in range(len(python_response.parameter_names)):
-            planned_parameter = plan_pb2.PlannedParameter(parameter_value=ares_value_utils.create_ares_value(python_response.parameter_values[i]))
-            planned_parameter.parameter_name = python_response.parameter_names[i]
-            new_planned_parameter = response_proto.planned_parameters.add()
-            new_planned_parameter.CopyFrom(planned_parameter)
+        elif isinstance(python_response, List) and all(isinstance(item, Plan) for item in python_response):
+            response_proto.plans.extend(plan_response_utils.python_plan_to_proto_plan(p) for p in python_response)
 
+        else:
+            response_proto.error_string = "The returned response from the user planning method was not valid, users must return either a list of plans or a plan response."
+            response_proto.planning_outcome = ares_outcome_enum_pb2.FAILURE
+
         print("Sending Plan Response.....")
         return response_proto
 

PyAres/Utils/ares_plan_status_code_utils.py

@@ -0,0 +1,10 @@
+from ..Models import PlanStatusCode
+from ares_datamodel.planning import plan_pb2
+from typing import cast
+
+def python_plan_status_to_proto_plan_status(py_value: PlanStatusCode) -> plan_pb2.PlanStatusCode:
+  val = cast(plan_pb2.PlanStatusCode, py_value)
+  return val
+
+def proto_plan_status_to_python_plan_status(proto_value: plan_pb2.PlanStatusCode) -> PlanStatusCode:
+  return PlanStatusCode(proto_value)

PyAres/Utils/plan_response_utils.py

@@ -1,6 +1,6 @@
 from ares_datamodel.planning import plan_pb2
-from ..Planning import PlanResponse
-from . import ares_outcome_utils
+from ..Planning import PlanResponse, Plan
+from . import ares_outcome_utils, planning_param_utils
 
 def proto_plan_response_to_python(proto_response: plan_pb2.PlanningResponse) -> PlanResponse:
   param_names = []
@@ -14,3 +14,11 @@ def proto_plan_response_to_python(proto_response: plan_pb2.PlanningResponse) ->
   python_response.outcome = ares_outcome_utils.proto_ares_outcome_to_python_ares_outcome(proto_response.planning_outcome)
   python_response.error_string = proto_response.error_string
   return python_response
+
+def python_plan_to_proto_plan(python_plan: Plan) -> plan_pb2.Plan:
+  proto_plan = plan_pb2.Plan()
+  proto_plan.planned_parameters.extend(planning_param_utils.convert_python_planned_param_to_proto_planned_param(p) for p in python_plan.planned_parameters)
+  proto_plan.planning_outcome = ares_outcome_utils.python_ares_outcome_to_proto_ares_outcome(python_plan.outcome)
+  proto_plan.error_string = python_plan.error_string
+
+  return proto_plan

PyAres/Utils/planning_param_utils.py

@@ -1,6 +1,5 @@
-from ..Planning import PlanningParameter, PlanResponse
+from ..Planning import PlanningParameter, PlannedParameter
 from ares_datamodel.planning import plan_pb2
-from typing import List
 from . import param_history_info_utils
 from . import ares_data_type_utils
 from . import ares_value_utils
@@ -17,4 +16,29 @@ def convert_python_plan_param_to_proto(param: PlanningParameter) -> plan_pb2.Pla
   new_proto_param.planner_name = param.planner_name
   ares_value_utils.py_to_ares_value(param.initial_value, new_proto_param.initial_value)
 
-  return new_proto_param 
+  return new_proto_param 
+
+def convert_proto_plan_param_to_python(param: plan_pb2.PlanningParameter) -> PlanningParameter:
+  new_python_param = PlanningParameter()
+  new_python_param.name = param.parameter_name
+  new_python_param.minimum_value = param.minimum_value
+  new_python_param.maximum_value = param.maximum_value
+  new_python_param.param_history.extend([param_history_info_utils.convert_proto_param_history_to_python(p) for p in param.parameter_history])
+  new_python_param.data_type = ares_data_type_utils.proto_ares_type_to_python_ares_type(param.data_type)
+  new_python_param.is_planned = param.is_planned
+  new_python_param.is_result = param.is_result
+  new_python_param.planner_name = param.planner_name
+  ares_value_utils.ares_value_to_py(param.initial_value, new_python_param.initial_value)
+
+  return new_python_param
+
+def convert_python_planned_param_to_proto_planned_param(param: PlannedParameter) -> plan_pb2.PlannedParameter:
+  new_proto_param = plan_pb2.PlannedParameter()
+  new_proto_param.parameter_name = param.parameter_name
+  ares_value_utils.py_to_ares_value(param.parameter_value, new_proto_param.parameter_value)
+
+  return new_proto_param
+
+def convert_proto_planned_param_to_proto_planned_param(param: plan_pb2.PlannedParameter) -> PlannedParameter:
+  new_python_param = PlannedParameter(parameter_name=param.parameter_name, parameter_value=ares_value_utils.ares_value_to_py(param.parameter_value))
+  return new_python_param

PyAres/__init__.py

@@ -3,6 +3,8 @@
 from .Planning import PlanRequest
 from .Planning import PlanningParameter
 from .Planning import ParameterHistoryItem
+from .Planning import PlannedParameter
+from .Planning import Plan
 from .Analyzing import AresAnalyzerService
 from .Analyzing import AnalysisResponse
 from .Analyzing import AnalysisRequest
@@ -17,4 +19,5 @@
 from .Models import AresSchemaEntry
 from .Models import Quantity
 from .Models import QuantitySchema
-from .Models import Limits
+from .Models import Limits
+from .Models import PlanStatusCode

pyproject.toml

@@ -16,7 +16,7 @@ classifiers = [
 ]
 
 dependencies = [
-    "ares-datamodel >= 0.29.0",
+    "ares-datamodel >= 0.30.0",
 ]
 
 [project.urls]