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/__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..c26d2fb 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,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__() \ 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_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 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..aeb7c69 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() + 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 \ 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/pyproject.toml b/pyproject.toml index 8b60d32..c88e589 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ classifiers = [ ] dependencies = [ - "ares-datamodel >= 0.29.0", + "ares-datamodel >= 0.30.0", ] [project.urls]