diff --git a/app/constants/websockets_constants.py b/app/constants/websockets_constants.py index c6d16652..43890503 100644 --- a/app/constants/websockets_constants.py +++ b/app/constants/websockets_constants.py @@ -25,6 +25,8 @@ # Enum Keys for different types of messages currently supported by the tool class MessageTypeEnum(str, Enum): PROMPT_REQUEST = "prompt_request" + OPTIONS_REQUEST = "options_request" + MESSAGE_REQUEST = "message_request" PROMPT_RESPONSE = "prompt_response" TEST_UPDATE = "test_update" TIME_OUT_NOTIFICATION = "time_out_notification" diff --git a/app/tests/user_prompts/test_options_select_user_prompts.py b/app/tests/user_prompts/test_options_select_user_prompts.py index 4034edda..666640be 100644 --- a/app/tests/user_prompts/test_options_select_user_prompts.py +++ b/app/tests/user_prompts/test_options_select_user_prompts.py @@ -298,7 +298,7 @@ def __expected_response_dict_cancelled() -> Dict[str, Any]: def __expected_prompt_dict(timeout: int = 2, message_id: int = 1) -> Dict[str, Any]: return { - MessageKeysEnum.TYPE: MessageTypeEnum.PROMPT_REQUEST, + MessageKeysEnum.TYPE: MessageTypeEnum.OPTIONS_REQUEST, MessageKeysEnum.PAYLOAD: { OPTIONS_KEY: test_options, TIMEOUT_KEY: timeout, diff --git a/app/user_prompt_support/user_prompt_manager.py b/app/user_prompt_support/user_prompt_manager.py index ba152fae..fbb0e48b 100644 --- a/app/user_prompt_support/user_prompt_manager.py +++ b/app/user_prompt_support/user_prompt_manager.py @@ -50,8 +50,16 @@ def __init__(self, prompt: PromptRequest, message_id: int) -> None: def as_dictionary(self) -> Dict[MessageKeysEnum, Any]: prompt_dict = self.prompt.dict() prompt_dict[MESSAGE_ID_KEY] = self.message_id + message_type = MessageTypeEnum.PROMPT_REQUEST + + # Check if options exist + if "options" in prompt_dict: + if prompt_dict["options"]: + message_type = MessageTypeEnum.OPTIONS_REQUEST + else: + message_type = MessageTypeEnum.MESSAGE_REQUEST message_dict = { - MessageKeysEnum.TYPE: MessageTypeEnum.PROMPT_REQUEST, + MessageKeysEnum.TYPE: message_type, MessageKeysEnum.PAYLOAD: prompt_dict, } return message_dict