Prompt request refactor

app/constants/websockets_constants.py

@@ -27,6 +27,7 @@ class MessageTypeEnum(str, Enum):
     PROMPT_REQUEST = "prompt_request"
     OPTIONS_REQUEST = "options_request"
     MESSAGE_REQUEST = "message_request"
+    FILE_UPLOAD_REQUEST = "file_upload_request"
     PROMPT_RESPONSE = "prompt_response"
     TEST_UPDATE = "test_update"
     TIME_OUT_NOTIFICATION = "time_out_notification"

app/user_prompt_support/prompt_request.py

@@ -13,43 +13,45 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from enum import Enum
 from typing import Dict, Optional
 
 from pydantic import BaseModel
 
-default_timeout_s = 60  # Seconds
-
+from app.constants.websockets_constants import MessageTypeEnum
 
-class PromptRequestType(str, Enum):
-    BASE = "base"
-    OPTIONS = "options"
-    TEXT = "text"
-    FILE = "file"
+default_timeout_s = 60  # Seconds
 
 
 class PromptRequest(BaseModel):
     prompt: Optional[str]
     timeout: int = default_timeout_s
-    __type: PromptRequestType = PromptRequestType.BASE
 
     @property
-    def type(self) -> PromptRequestType:
-        return self.__type
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.INVALID_MESSAGE
 
 
 class OptionsSelectPromptRequest(PromptRequest):
-    __type = PromptRequestType.OPTIONS
     options: Dict[str, int]
 
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.OPTIONS_REQUEST
+
 
 class TextInputPromptRequest(PromptRequest):
-    __type = PromptRequestType.TEXT
     placeholder_text: Optional[str]
     default_value: Optional[str]
     regex_pattern: Optional[str]
 
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.PROMPT_REQUEST
+
 
 class UploadFilePromptRequest(PromptRequest):
-    __type = PromptRequestType.FILE
     path: str = "api/v1/test_run_execution/file_upload/"
+
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.FILE_UPLOAD_REQUEST

app/user_prompt_support/user_prompt_manager.py

@@ -50,16 +50,9 @@ 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: message_type,
+            MessageKeysEnum.TYPE: self.prompt.messageType,
             MessageKeysEnum.PAYLOAD: prompt_dict,
         }
         return message_dict
@@ -227,4 +220,4 @@ def __prompt_exchange_for_message(
         )
 
 
-user_prompt_manager = UserPromptManager()
+user_prompt_manager: UserPromptManager = UserPromptManager()