diff --git a/chatmastermind/message.py b/chatmastermind/message.py index 443455e..9623195 100644 --- a/chatmastermind/message.py +++ b/chatmastermind/message.py @@ -92,7 +92,7 @@ class MessageFilter: class AILine(str): """ - A line that represents the AI name in a '.txt' file.. + A line that represents the AI name in the 'txt' format. """ prefix: Final[str] = 'AI:' @@ -112,7 +112,7 @@ class AILine(str): class ModelLine(str): """ - A line that represents the model name in a '.txt' file.. + A line that represents the model name in the 'txt' format. """ prefix: Final[str] = 'MODEL:' @@ -216,7 +216,10 @@ class Message(): model: Optional[str] = field(default=None, compare=False) file_path: Optional[pathlib.Path] = field(default=None, compare=False) # class variables - file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml'] + file_suffixes_read: ClassVar[list[str]] = ['.msg', '.txt', '.yaml'] + file_suffix_write: ClassVar[str] = '.msg' + file_formats: ClassVar[list[str]] = ['txt', 'yaml'] + default_format: ClassVar[str] = 'txt' tags_yaml_key: ClassVar[str] = 'tags' file_yaml_key: ClassVar[str] = 'file_path' ai_yaml_key: ClassVar[str] = 'ai' @@ -276,24 +279,16 @@ class Message(): tags: set[Tag] = set() if not file_path.exists(): raise MessageError(f"Message file '{file_path}' does not exist") - if file_path.suffix not in cls.file_suffixes: + if file_path.suffix not in cls.file_suffixes_read: raise MessageError(f"File type '{file_path.suffix}' is not supported") - # for TXT, it's enough to read the TagLine - if file_path.suffix == '.txt': - with open(file_path, "r") as fd: - try: - tags = TagLine(fd.readline()).tags(prefix, contain) - except TagError: - pass # message without tags - else: # '.yaml' - try: - message = cls.from_file(file_path) - if message: - msg_tags = message.filter_tags(prefix=prefix, contain=contain) - except MessageError as e: - print(f"Error processing message in '{file_path}': {str(e)}") - if msg_tags: - tags = msg_tags + try: + message = cls.from_file(file_path) + if message: + msg_tags = message.filter_tags(prefix=prefix, contain=contain) + except MessageError as e: + print(f"Error processing message in '{file_path}': {str(e)}") + if msg_tags: + tags = msg_tags return tags @classmethod @@ -328,15 +323,16 @@ class Message(): """ if not file_path.exists(): raise MessageError(f"Message file '{file_path}' does not exist") - if file_path.suffix not in cls.file_suffixes: + if file_path.suffix not in cls.file_suffixes_read: raise MessageError(f"File type '{file_path.suffix}' is not supported") - - if file_path.suffix == '.txt': + # try TXT first + try: message = cls.__from_file_txt(file_path, mfilter.tags_or if mfilter else None, mfilter.tags_and if mfilter else None, mfilter.tags_not if mfilter else None) - else: + # then YAML + except MessageError: message = cls.__from_file_yaml(file_path) if message and (mfilter is None or message.match(mfilter)): return message @@ -442,21 +438,25 @@ class Message(): output.append(self.answer) return '\n'.join(output) - def to_file(self, file_path: Optional[pathlib.Path]=None) -> None: # noqa: 11 + def to_file(self, file_path: Optional[pathlib.Path]=None, mformat: str = Message.default_format) -> None: # noqa: 11 """ - Write a Message to the given file. Type is determined based on the suffix. - Currently supported suffixes: ['.txt', '.yaml'] + Write a Message to the given file. Supported message file formats are 'txt' and 'yaml'. + Suffix is always '.msg'. """ if file_path: self.file_path = file_path if not self.file_path: raise MessageError("Got no valid path to write message") - if self.file_path.suffix not in self.file_suffixes: - raise MessageError(f"File type '{self.file_path.suffix}' is not supported") + if mformat not in self.file_formats: + raise MessageError(f"File format '{mformat}' is not supported") + # do not write old suffixes + if self.file_path.suffix != self.file_suffix_write: + raise MessageError(f"File suffix '{self.file_path.suffix}' is not supported") # TXT - if self.file_path.suffix == '.txt': + if mformat == 'txt': return self.__to_file_txt(self.file_path) - elif self.file_path.suffix == '.yaml': + # YAML + elif mformat == 'yaml': return self.__to_file_yaml(self.file_path) def __to_file_txt(self, file_path: pathlib.Path) -> None: @@ -468,8 +468,8 @@ class Message(): * Model [Optional] * Question.txt_header * Question - * Answer.txt_header - * Answer + * Answer.txt_header [Optional] + * Answer [Optional] """ with tempfile.NamedTemporaryFile(dir=file_path.parent, prefix=file_path.name, mode="w", delete=False) as temp_fd: temp_file_path = pathlib.Path(temp_fd.name)