add supported method to converter

Azure · Apr 9, 2024 · 87476e5 · 87476e5
1 parent 9564e59
commit 87476e5
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/pyrit/prompt_converter/file_converter/txt_file_converter.py b/pyrit/prompt_converter/file_converter/txt_file_converter.py
@@ -27,8 +27,9 @@ def convert(self, *, prompt: str, input_type: PromptDataType = "text") -> str:
         Returns:
             str: The file name.
         """
-        if input_type != "text":
-            raise ValueError("TxtFileConverter only supports input type 'text'.")
+        if not self.is_supported(input_type):
+            raise ValueError("Input type not supported")
+
         if not self._file_name:
             self._file_name = uuid4()
         full_file_name = f"{self._file_name}.txt"
@@ -38,3 +39,7 @@ def convert(self, *, prompt: str, input_type: PromptDataType = "text") -> str:
             "data": prompt,
             "file_location": full_file_name
         })
+
+
+    def is_supported(self, input_type: PromptDataType) -> bool:
+        return input_type == "text"