diff --git a/PyAres/Demo/Planners/planner_wiki.py b/PyAres/Demo/Planners/planner_wiki.py index 8315230..12b5f80 100644 --- a/PyAres/Demo/Planners/planner_wiki.py +++ b/PyAres/Demo/Planners/planner_wiki.py @@ -1,23 +1,27 @@ -from PyAres import AresPlannerService, PlanRequest, PlanResponse, AresDataType +from PyAres import * +from typing import List import random -def generate_plan(request: PlanRequest) -> PlanResponse: - planned_values = [] - names = [] +def generate_plans(request: PlanRequest) -> List[Plan]: + plans = [] + for i in request.previous_plan_status_codes: + print(f"Status Code: {i}") - # Iterate through every parameter configured in the ARES Experiment - for param in request.parameters: - # Simple Logic: Pick a random value within the allowed range - val = random.uniform(param.minimum_value, param.maximum_value) + # Make as many plans as was requested + for i in range(request.batch_size): + current_params = [] + for param in request.parameters: + # Simple Logic: Pick a random value within the allowed range + val = random.uniform(param.minimum_value, param.maximum_value) + current_params.append(PlannedParameter(param.name, val)) - names.append(param.name) - planned_values.append(val) + plans.append(Plan(current_params, Outcome.SUCCESS)) - return PlanResponse(parameter_names=names, parameter_values=planned_values) + return plans if __name__ == "__main__": service = AresPlannerService( - generate_plan, + generate_plans, "Random Search Planner", "This planner picks random values within bounds.", "1.0.0" 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 34cf4df..2ba85d3 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/__init__.py b/PyAres/Planning/__init__.py index cbd1ff6..4fa1897 100644 --- a/PyAres/Planning/__init__.py +++ b/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", ] diff --git a/PyAres/Planning/planner_models.py b/PyAres/Planning/planner_models.py index fd7ced9..129bc1b 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 """ @@ -28,9 +29,15 @@ class PlanningParameter: Designed to provide a more user-friendly abstraction for the user to interact with planning parameters through. """ - def __init__(self, name: str, minimum_value: float, - maximum_value: float, param_history: list[ParameterHistoryItem], data_type: AresDataType, - is_planned: bool, is_result: bool, planner_name: str, initial_value = None): + def __init__(self, name: str, + data_type: AresDataType, + minimum_value: float = None, + maximum_value: float = None, + param_history: list[ParameterHistoryItem] = [], + is_planned: bool = False, + is_result: bool = False, + planner_name: str = "", + initial_value = None): """ Initializes a PlanningParameter. @@ -107,7 +114,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], + metadata: RequestMetadata = RequestMetadata.from_default_values(), + batch_size: int = 1, + previous_plan_status_codes: List[PlanStatusCode] = None): """ Initializes a PlanRequest. @@ -117,8 +130,14 @@ 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 + if previous_plan_status_codes is None: + self.previous_plan_status_codes = [] + else: + self.previous_plan_status_codes = previous_plan_status_codes + def __str__(self) -> str: param_str = "\n ".join(self.parameter_names) settings_str = "\n ".join([f"{k}: {v}" for k, v in self.settings.items()]) @@ -127,13 +146,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 +196,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 +223,39 @@ def __str__(self): f" error_string: {self.error_string}\n") def __repr__(self) -> str: - return self.__str__() \ No newline at end of file + 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})" diff --git a/PyAres/Planning/planning_service.py b/PyAres/Planning/planning_service.py index 11e938d..ba95853 100644 --- a/PyAres/Planning/planning_service.py +++ b/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 diff --git a/PyAres/Utils/ares_outcome_utils.py b/PyAres/Utils/ares_outcome_utils.py index 930b881..1de3328 100644 --- a/PyAres/Utils/ares_outcome_utils.py +++ b/PyAres/Utils/ares_outcome_utils.py @@ -2,10 +2,9 @@ from ares_datamodel import ares_outcome_enum_pb2 from typing import cast -def python_ares_outcome_to_proto_ares_outcome(py_value: Outcome) -> ares_outcome_enum_pb2.Outcome: +def python_ares_outcome_to_proto_ares_outcome(py_value: Outcome) -> int: """ A method to convert from the python AresDataType class to the protobuf version """ - val = cast(ares_outcome_enum_pb2.Outcome, py_value.value) - return val + return py_value.value def proto_ares_outcome_to_python_ares_outcome(proto_value: ares_outcome_enum_pb2.Outcome) -> Outcome: """ A method to convert from the protobuf AresDataType class to the python version """ 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..9c74d38 --- /dev/null +++ b/PyAres/Utils/ares_plan_status_code_utils.py @@ -0,0 +1,9 @@ +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) -> int: + return py_value.value + +def proto_plan_status_to_python_plan_status(proto_value: plan_pb2.PlanStatusCode) -> PlanStatusCode: + return PlanStatusCode(proto_value) \ No newline at end of file diff --git a/PyAres/Utils/ares_value_utils.py b/PyAres/Utils/ares_value_utils.py index 0af2c75..c26f6fb 100644 --- a/PyAres/Utils/ares_value_utils.py +++ b/PyAres/Utils/ares_value_utils.py @@ -7,7 +7,7 @@ from . import ares_data_type_utils from ..Models import AresDataType, Quantity -def ares_value_to_py(ares_value: ares_struct_pb2.AresValue): +def ares_value_to_py(ares_value: ares_struct_pb2.AresValue) -> Any: """Converts an AresValue protobuf message to a Python native type.""" field = ares_value.WhichOneof("kind") if field == "null_value": diff --git a/PyAres/Utils/device_status_code_utils.py b/PyAres/Utils/device_status_code_utils.py index fbcde35..4be720a 100644 --- a/PyAres/Utils/device_status_code_utils.py +++ b/PyAres/Utils/device_status_code_utils.py @@ -2,10 +2,9 @@ from ares_datamodel import command_status_code_pb2 from typing import cast -def python_status_code_to_proto_status_code(py_value: StatusCode) -> command_status_code_pb2.CommandStatusCode: +def python_status_code_to_proto_status_code(py_value: StatusCode) -> int: """ A method to convert from the python StatusCode enum class to the protobuf version """ - val = cast( command_status_code_pb2.CommandStatusCode, py_value.value) - return val + return py_value.value def proto_status_code_to_python_status_code(proto_value: command_status_code_pb2.CommandStatusCode) -> StatusCode: """ A method to convert from the protobuf StatusCode enum class to the python version """ diff --git a/PyAres/Utils/plan_response_utils.py b/PyAres/Utils/plan_response_utils.py index 7b237fd..d60f670 100644 --- a/PyAres/Utils/plan_response_utils.py +++ b/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 \ No newline at end of file diff --git a/PyAres/Utils/planning_param_utils.py b/PyAres/Utils/planning_param_utils.py index 1859da0..026cc52 100644 --- a/PyAres/Utils/planning_param_utils.py +++ b/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 \ No newline at end of file + return new_proto_param + +def convert_proto_plan_param_to_python(param: plan_pb2.PlanningParameter) -> PlanningParameter: + new_python_param = PlanningParameter(param.parameter_name, ares_data_type_utils.proto_ares_type_to_python_ares_type(param.data_type)) + 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.is_planned = param.is_planned + new_python_param.is_result = param.is_result + new_python_param.planner_name = param.planner_name + + if(param.initial_value is not None): + new_python_param.initial_value = ares_value_utils.ares_value_to_py(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 \ No newline at end of file diff --git a/PyAres/__init__.py b/PyAres/__init__.py index 1c51655..413969d 100644 --- a/PyAres/__init__.py +++ b/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 \ No newline at end of file +from .Models import Limits +from .Models import PlanStatusCode \ No newline at end of file diff --git a/tests/test_batch_planning_regressions.py b/tests/test_batch_planning_regressions.py new file mode 100644 index 0000000..1c2aa90 --- /dev/null +++ b/tests/test_batch_planning_regressions.py @@ -0,0 +1,64 @@ +import unittest + +from ares_datamodel import ares_data_type_pb2 +from ares_datamodel.planning import plan_pb2 + +from PyAres.Models import AresDataType, PlanStatusCode, RequestMetadata +from PyAres.Planning import PlanRequest +from PyAres.Utils import ares_plan_status_code_utils, planning_param_utils + + +class TestBatchPlanningRegressions(unittest.TestCase): + def test_convert_proto_planning_parameter_to_python(self): + proto_parameter = plan_pb2.PlanningParameter( + parameter_name="temperature", + minimum_value=20.0, + maximum_value=100.0, + data_type=ares_data_type_pb2.AresDataType.NUMBER, + is_planned=True, + is_result=False, + planner_name="Test Planner", + initial_value=None) + proto_parameter.initial_value.number_value = 25.0 + + parameter = planning_param_utils.convert_proto_plan_param_to_python(proto_parameter) + + self.assertEqual(parameter.name, "temperature") + self.assertEqual(parameter.minimum_value, 20.0) + self.assertEqual(parameter.maximum_value, 100.0) + self.assertEqual(parameter.data_type, AresDataType.NUMBER) + self.assertTrue(parameter.is_planned) + self.assertFalse(parameter.is_result) + self.assertEqual(parameter.planner_name, "Test Planner") + self.assertEqual(parameter.initial_value, 25.0) + + def test_python_status_can_be_added_to_protobuf_request(self): + status = ares_plan_status_code_utils.python_plan_status_to_proto_plan_status(PlanStatusCode.PLAN_ACCEPTED) + request = plan_pb2.PlanningRequest() + + request.previous_plan_status_codes.append(status) + + self.assertEqual( + request.previous_plan_status_codes[0], + plan_pb2.PLAN_ACCEPTED, + ) + + def test_fourth_positional_argument_remains_metadata(self): + metadata = RequestMetadata.from_default_values() + + request = PlanRequest([], {}, [], metadata) + + self.assertIs(request.request_metadata, metadata) + self.assertEqual(request.batch_size, 1) + + def test_default_status_code_lists_are_independent(self): + first_request = PlanRequest([], {}, []) + second_request = PlanRequest([], {}, []) + + first_request.previous_plan_status_codes.append(PlanStatusCode.PLAN_ACCEPTED) + + self.assertEqual(second_request.previous_plan_status_codes, []) + + +if __name__ == "__main__": + unittest.main(verbosity=2)