From c1811664420e549040af23d2b8b2ffaba2f46cff Mon Sep 17 00:00:00 2001 From: Nick Kleiner Date: Thu, 18 Jun 2026 13:30:29 -0400 Subject: [PATCH] Adding batch planning capabilities to PyAres --- PyAres/Models/__init__.py | 5 +++-- PyAres/Models/ares_data_models.py | 6 ++++++ PyAres/Planning/planner_models.py | 13 +++++++++++-- PyAres/Planning/planning_service.py | 4 +++- PyAres/Utils/ares_plan_status_code_utils.py | 10 ++++++++++ 5 files changed, 33 insertions(+), 5 deletions(-) create mode 100644 PyAres/Utils/ares_plan_status_code_utils.py diff --git a/PyAres/Models/__init__.py b/PyAres/Models/__init__.py index 342beba..0d86d75 100644 --- a/PyAres/Models/__init__.py +++ b/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" ] \ No newline at end of file diff --git a/PyAres/Models/ares_data_models.py b/PyAres/Models/ares_data_models.py index 6a06019..613d3bf 100644 --- a/PyAres/Models/ares_data_models.py +++ b/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 diff --git a/PyAres/Planning/planner_models.py b/PyAres/Planning/planner_models.py index fd7ced9..b12dfde 100644 --- a/PyAres/Planning/planner_models.py +++ b/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,12 @@ 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], + metadata: RequestMetadata = RequestMetadata.from_default_values(), + previous_plan_status_code: PlanStatusCode = PlanStatusCode.PLAN_STATUS_UNSPECIFIED): """ Initializes a PlanRequest. @@ -118,6 +124,7 @@ def __init__(self, parameters: list[PlanningParameter], settings: Dict[str, Any] self.settings = settings self.analysis_results = analysis_results self.request_metadata = metadata + self.previous_plan_status_code = previous_plan_status_code def __str__(self) -> str: param_str = "\n ".join(self.parameter_names) @@ -175,6 +182,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()) diff --git a/PyAres/Planning/planning_service.py b/PyAres/Planning/planning_service.py index 11e938d..4e96211 100644 --- a/PyAres/Planning/planning_service.py +++ b/PyAres/Planning/planning_service.py @@ -21,6 +21,7 @@ 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 # Import python models from ..Models import ares_data_models, Limits @@ -122,7 +123,8 @@ 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), + previous_plan_status_code=ares_plan_status_code_utils.proto_plan_status_to_python_plan_status(request.previous_plan_status_code)) #Handle call using the user's custom planning logic response_proto = plan_pb2.PlanningResponse() diff --git a/PyAres/Utils/ares_plan_status_code_utils.py b/PyAres/Utils/ares_plan_status_code_utils.py new file mode 100644 index 0000000..6f93dfc --- /dev/null +++ b/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) \ No newline at end of file